from openai import OpenAI
import urllib
import requests
import base64
import os
import ast
import cv2
from io import BytesIO
from PIL import Image
from tempfile import NamedTemporaryFile
import pyheif
import time
from zipfile import ZipFile
import gradio as gr
from docx import Document
import numpy as np


api_key = os.environ['OPENAI_API_KEY']
brandfolder_api = os.environ['BRANDFOLDER_API_KEY']

client_key_dict = {
    "The Official Moving Company, LLC": 'KXRbpext',
    "Newmark Commercial Real Estate": 'none',
    "Test Collection": 'test',
    'Direct Mail Xperts LLC':'d5J3MdlO'
}

section_key_dict = {
    "Original Project Assets": 'c5vm8cnh9jvkjbh7r43qxkv',
    "Pre-Processed Images": 'rfqf67pbhn8hg6pjcj762q3q',
    "AI Processed Images": 'czpq4nwz78c3cwnp6h9n44z'
}

# Functions

def rename(filename):
    client = OpenAI()
    completion = client.chat.completions.create(
      model="gpt-4o",
      messages=[
        {"role": "system", "content": "You are a helpful assistant specializing in renaming files."},
        {"role": "user", "content": f"Provide a similar name for this filename: {filename}. Only return the filename and use hyphens in the filename."}
      ]
    )
    return completion.choices[0].message.content


def get_collection_dict():
  headers = {
  'Accept': 'application/json',
  'Authorization': brandfolder_api
  }

  r = requests.get('https://brandfolder.com/api/v4/brandfolders/988cgqcg8xsrr5g9h7gtsqkg/collections?per=300', params={
      # use a dict with your desired URL parameters here
  }, headers=headers)

  temp = r.json()['data']

  collection_dict = {item['attributes']['name']:item['id'] for item in temp}

  return collection_dict


def get_collection_names():
    collection_dict = get_collection_dict()
    return list(collection_dict.keys())


def get_topical_map_text(path):
    document = Document(path)

    extracted_text = []

    for paragraph in document.paragraphs:
        # Get the left indentation of the current paragraph (if any)
        left_indent = paragraph.paragraph_format.left_indent
        if left_indent == None:
          continue
        else:
          indent_level = int(left_indent.pt / 20)  # Convert Twips to points and then to a simple indentation level

          # You might want to adjust the logic below depending on how you want to represent indentation
          indent_symbol = " " * indent_level  # This creates a number of spaces based on the indentation level; adjust as needed

          # Construct the paragraph text with indentation representation
          formatted_text = f"{indent_symbol}{paragraph.text}"
          extracted_text.append(formatted_text)

    return "\n".join(extracted_text)


def get_asset_info(asset_id):
  '''
  Takes information from asset_id

  Input: asset_id
  Output: collection_id, collection_name, section_id
  '''
  # asset_id = data['data']['attributes']['key']
  headers = {
    'Content-Type': 'application/json',
    'Authorization': brandfolder_api
  }
  r = requests.get(f'https://brandfolder.com/api/v4/assets/{asset_id}?include=section,collections,custom_fields,attachments', params={}, headers=headers)

  # gets section_id
  try:
    section_id = r.json()['data']['relationships']['section']['data']['id']
  except:
    section_id = ''

  # gets collection_id
  # gets collection_name
  try:
    collection_id = r.json()['data']['relationships']['collections']['data'][0]['id']
    collection_name = [item['attributes']['name'] for item in r.json()['included'] if item['type']=='collections'][0]
  except:
    collection_id = ''
    collection_name = ''

  # gets asset_name, asset_type, and asset_url
  try:
    asset_type = [item['attributes']['value'] for item in r.json()['included'] if item['type'] == 'custom_field_values' and item['attributes']['value']=='Photo'][0]
  except:
    asset_type = ''
  try:
    asset_name = r.json()['data']['attributes']['name']
  except:
    asset_name = ''
  try:
    access_key = [item['attributes']['value'] for item in r.json()['included'] if item['type'] == 'custom_field_values' and item['attributes']['key'] == 'What is your Access Code?'][0]
  except:
    access_key = ''
  try:
    asset_url = [item['attributes']['url'] for item in r.json()['included'] if item['type'] == 'attachments'][0]
  except:
    asset_url = ''
  try:
    client_name = [item['attributes']['value'] for item in r.json()['included'] if item['type'] == 'custom_field_values' and item['attributes']['key'] == 'Client Name'][0]
  except:
    client_name = ''
  try:
    project_name = [item['attributes']['value'] for item in r.json()['included'] if item['type'] == 'custom_field_values' and item['attributes']['key'] == 'List Project Name Photos Belong To'][0]
  except:
    project_name = ''

  return_dict = {
      "section_id": section_id,
      "collection_id": collection_id,
      "collection_name": collection_name,
      "asset_type": asset_type,
      "asset_name": asset_name,
      "access_key": access_key,
      "image_url": asset_url,
      "client_name": client_name,
      "project_name": project_name
  }

  return return_dict


def get_seo_tags(image_url, topical_map, new_imgs, attempts=0, max_attempts=5):
  '''
  Gets the seo tags and topic/sub-topic classification for an image using OpenAI GPT-4 Vision Preview
  Input: image path of desired file
  Output: dict of topic, sub-topic, and seo tags
  '''

  if attempts > max_attempts:
    print("Maximum number of retries exceeded.")
    return {"error": "Max retries exceeded, operation failed."}
      
  print('in seo_tags')
  
  # Query for GPT-4
  topic_map_query = f"""
  % You are an expert web designer that can only answer questions relevent to the following Topical Map.
  % Goal: Output the topic, description, caption, seo tags, alt_tags, and filename for this image using the Topical Map provided.
  
  % TOPCIAL MAP
  ```{topical_map}```

  """
 # IF YOU CANNOT PROVIDE AN TOPIC FOR EVERY IMAGE AFTER 5 ATTEMPTS, REPLY WITH 'irrelevant'.
  topic_list = topical_map.split('\n')
  topic_list = [topic.strip() for topic in topic_list]
  topic_list.insert(0, "irrelevant")

  
  def compress_and_encode_image(url, target_size_mb=20, quality=70):
    # Fetch the image from the URL
    response = requests.get(url, stream=True)
    response.raise_for_status()

    # Open the image using Pillow, handling HEIC files
    img = None
    img_format = response.headers['Content-Type'].split('/')[-1]
    
    if img_format.lower() == 'heic':
        heif_file = pyheif.read_heif(response.content)
        img = Image.frombytes(heif_file.mode, heif_file.size, heif_file.data, "raw", heif_file.mode, heif_file.stride)
    else:
        img = Image.open(BytesIO(response.content))

    img = img.convert('RGB')

    # Compress the image by adjusting the quality
    img_bytes = BytesIO()
    img.save(img_bytes, format='JPEG', quality=quality)

    # Check if the image size is acceptable
    while img_bytes.getbuffer().nbytes > (target_size_mb * 1024 * 1024) and quality > 10:
        quality -= 5
        img_bytes = BytesIO()
        img.save(img_bytes, format='JPEG', quality=quality)

    # Encode the image content to base64
    encoded_image = base64.b64encode(img_bytes.getvalue()).decode('utf-8')

    return encoded_image

  base64_image = compress_and_encode_image(image_url)

  # REMOVE WHEN SHARING FILE
  api_key = os.environ['OPENAI_API_KEY']
  
  # Calling gpt-4 vision
  headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {api_key}"
  }
# IF YOU CANNOT PROVIDE AN TOPIC FOR EVERY IMAGE AFTER 5 ATTEMPTS, REPLY WITH 'irrelevant'.
  payload = {
    "model": "gpt-4o",
    "response_format": {"type": "json_object"},
    "messages": [
      {'role': 'system', 'content': 'You are an expert web designer that can only answer questions relevent to the following topical map.'
      },
      {
        "role": "user",
        "content": [
          {
            "type": "text",
            "text": topic_map_query +
                    """
                    % INSTRUCTIONS
                        Step 1 - Generate keywords to describe this image
                        Step 2 - Decide which topic in the Topicla Map this image fall under, using the keywords you generated and the image itself. You are only permitted to use the exact wording of the topic in the topical map. 
                        Step 2 - Provide a topic-relevant 5 sentence description for the image. Describe the image only using context relevant to the topics in the topical map. 
                                        Adhere to the following guidelines when crafting your 5 sentence description: 
                                        - Mention only the contents of the image. 
                                        - Do not mention the quality of the image. 
                                        - Ignore all personal information within the image. 
                                        - Be as specific as possible when identifying tools/items in the image.
                        Step 3 - Using the description in Step 1, create a 160 character caption. Make sure the caption is less than 160 characters.
                        Step 4 - Using the description in Step 1, create 3 topic-relevant SEO tags for this image that will drive traffic to our website. The SEO tags must be two words or less. You must give 3 SEO tags.
                        Step 5 - Using the description in Step 1, provide a topic-relevant SEO alt tag for the image that will enhance how the website is ranked on search engines.
                        Step 6 - Using the description in Step 1, provide a new and unique filename for the image as well. Use hyphens for the filename. Do not include extension.
                        Step 7 - YOU ARE ONLY PERMITTED TO OUTPUT THE TOPIC, DESCRIPTION, CAPTION, SEO, ALT_TAG, AND FILENAME IN THE FOLLOWING JSON FORMAT:
                    
                    % OUTPUT FORMAT:
                        {"topic": topic,
                        "description": description,
                        "caption": caption,
                        "seo": [seo],
                        "alt_tag": [alt tag],
                        "filename": filename
                        }
                    """
          },
          {
            "type": "image_url",
            "image_url": {
              "url": f"data:image/jpeg;base64, {base64_image}"
            }
          }
        ]
      }
    ],
    "max_tokens": 300
  }

  try:
    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
    response_data = response.json()
    if response.status_code == 200 and 'choices' in response_data and len(response_data['choices']) > 0:
        keys = ['topic', 'description', 'caption', 'seo', 'alt_tag', 'filename']
        json_dict = ast.literal_eval(response.json()['choices'][0]['message']['content'])
        if json_dict['topic'] not in topic_list:
            return get_seo_tags(image_path, topical_map, new_imgs, attempts=attempts+1)
        if set(json_dict.keys()) != set(keys):
            return get_seo_tags(image_path, topical_map, new_imgs, attempts=attempts+1)
        
        
        return json_dict
    else:
        print("API call failed or bad data, retrying...")
        return get_seo_tags(image_url, topical_map, new_imgs, attempts=attempts + 1)
  except Exception as e:
    print("Exception during API call:", str(e))
    return get_seo_tags(image_url, topical_map, new_imgs, attempts=attempts + 1)

def personalize_answer(answer, query_engine):
    if query_engine:
        prompt = f'''
        % You are an expert construction contracter describing this image
    
        % Goal: Add relevant company information to the Answer based on the context provided. Only output the Enhancement. Remove the triple quotes from the output
    
        % Answer:
        ```{answer}```
    
        % Instructions:
            Step 1 - Identify what relevant company information from the context is relevant to the Answer
            Step 2 - Enhance the Answer with the relevant company information. This will be known as the Enhancement
            Step 3 - Make the Enhancement the same character length as the Answer. Use Python to check that they are the same character length.
            Step 4 - Only output the Enhancement. Remove the triple quotes from the output
            
        '''
    
        response = query_engine.query(prompt)
    
        return response.response.replace("`", "")
    else:
        return answer

# creates the asset in the client's brand folder
def create_ai_asset(asset_dict, topical_map, collection_name, new_imgs, query_engine=None, tags=True):
    '''
    Creates asset from image path. Also creates seo tags, topic, and alt tag for
    image
    Input: name of initial asset, name of client, path to image, create tags boolean
    Output: id of asset
    '''
    print(asset_dict)
    # results from asset_dict
    topical_map = get_topical_map_text(topical_map)
    client_name = asset_dict['client_name']
    access_key = asset_dict['access_key']
    try:
        client_key = client_key_dict[collection_name]
    except:
        client_key = 'no key'
    if client_name != collection_name and access_key != client_key:
      print(f'{collection_name} != {client_name}')
      print(f'{access_key} != {client_key}')
      return 

    asset_name = asset_dict['asset_name']
    
    collection_id = asset_dict['collection_id']
    project_name = asset_dict['project_name']
    if collection_id == '':
      collection_dict_temp = get_collection_dict()
      collection_id = collection_dict_temp[client_name]
    image_url = asset_dict['image_url']

    # get seo, topic, and sub-topic from OpenAI API
    json_dict = get_seo_tags(image_url, topical_map, new_imgs)
    if not json_dict:
        json_dict = get_seo_tags(image_url, topical_map, new_imgs)

    # parsing out results from get_seo_tags
    topic = json_dict['topic']
    description = json_dict['description']
    caption = json_dict['caption']
    seo_tags = json_dict['seo']
    alt_tag = json_dict['alt_tag']
    image_name = json_dict['filename']

    description = personalize_answer(description, query_engine)
    caption = personalize_answer(caption, query_engine)
    alt_tag = personalize_answer(alt_tag, query_engine)
    
    headers = {
      'Content-Type': 'application/json',
      'Authorization': brandfolder_api
    }

    r = requests.get(f'https://brandfolder.com/api/v4/collections/{collection_id}/assets', params={
        # use a dict with your desired URL parameters here
    }, headers=headers)

    asset_names = [item['attributes']['name'] for item in r.json()['data']]

    asset_names = new_imgs + asset_names
    
    while image_name in asset_names:
        image_name = rename(image_name)
        

    # og image url
    og_object_url = image_url


    # binary upload of image_path
    r = requests.get('https://brandfolder.com/api/v4/upload_requests', params={}, headers=headers)

    # used to upload the image
    upload_url = r.json()['upload_url']
    # container for the uploaded image to be used by the post request
    object_url = r.json()['object_url']

    def download_and_resize_image(image_url, upload_url):
        # Fetch the image from the URL
        url_response = urllib.request.urlopen(image_url)
        img_array = np.array(bytearray(url_response.read()), dtype=np.uint8)
    
        # Try to decode the image using OpenCV
        image = cv2.imdecode(img_array, -1)
    
        # If the image is None, it might be a HEIC file
        if image is None:
            heif_file = pyheif.read_heif(img_array)
            img = Image.frombytes(
                heif_file.mode,
                heif_file.size,
                heif_file.data,
                "raw",
                heif_file.mode,
                heif_file.stride
            )
            # Convert to RGB
            img = img.convert('RGB')
            # Save to a BytesIO object
            img_bytes = BytesIO()
            img.save(img_bytes, format='JPEG')
            img_bytes.seek(0)
            img_array = np.array(bytearray(img_bytes.read()), dtype=np.uint8)
            # Decode the JPEG image using OpenCV
            image = cv2.imdecode(img_array, -1)
    
        # Resize the image based on its dimensions and area
        try:
            height, width, c = image.shape
        except:
            height, width = image.shape
    
        area = width * height
        if width > height:
            # Landscape image
            if area > 667000:
                image = cv2.resize(image, (1000, 667))
        else:
            # Portrait image
            if area > 442236:
                image = cv2.resize(image, (548, 807))
    
        # Save the image to a temporary file and upload it
        with NamedTemporaryFile(delete=True, suffix='.jpg') as temp_image:
            cv2.imwrite(temp_image.name, image)
            temp_image.seek(0)
            response = requests.put(upload_url, data=temp_image)
    
        return response

    response = download_and_resize_image(image_url, upload_url)
    
    # posts image with image name

    r = requests.post(f'https://brandfolder.com/api/v4/collections/{collection_id}/assets', json={
        # use a dict with the POST body here
        'data': {
            'attributes': [
                {
                    'name': image_name,
                    'description': description,
                    'attachments': [
                        {
                            'url': object_url,
                            'filename': f'{image_name}.jpg'
                        },
                        {
                            'url': og_object_url,
                            'filename': f'{image_name}-original.jpg'
                        }
                    ]
                }
            ]
        },
        # AI Original section key
        'section_key': 'czpq4nwz78c3cwnp6h9n44z'
      }, params={}, headers=headers)

    # id of newly created asset
    asset_id = r.json()['data'][0]['id']

    # tags and topic payloads
    tags_payload = {'data': {'attributes': [{'name': tag} for tag in seo_tags]}}


    topic_payload = {'data':
                [
                    {
                        'attributes': {
                            'value': topic
                            },
                            'relationships': {
                                'asset': {
                                    'data': {'type': 'assets', 'id': asset_id}
                                            }}
                    }]}


    alt_tag_payload = {'data':
                    [
                        {
                            'attributes': {
                                'value': alt_tag
                                },
                                'relationships': {
                                    'asset': {
                                        'data': {'type': 'assets', 'id': asset_id}
                                                }}
                        }]}

    year_payload = {'data':
                    [
                        {
                            'attributes': {
                                'value': 2024
                                },
                                'relationships': {
                                    'asset': {
                                        'data': {'type': 'assets', 'id': asset_id}
                                                }}
                        }]}

    client_payload = {'data':
                [
                    {
                        'attributes': {
                            'value': client_name
                            },
                            'relationships': {
                                'asset': {
                                    'data': {'type': 'assets', 'id': asset_id}
                                            }}
                    }]}
    caption_payload = {'data':
            [
                {
                    'attributes': {
                        'value': caption
                        },
                        'relationships': {
                            'asset': {
                                'data': {'type': 'assets', 'id': asset_id}
                                        }}
                }]}

    project_payload = {'data':
            [
                {
                    'attributes': {
                        'value': project_name
                        },
                        'relationships': {
                            'asset': {
                                'data': {'type': 'assets', 'id': asset_id}
                                        }}
                }]}

    year_id = 'k8vr5chnkw3nrnrpkh4f9fqm'
    client_name_id = 'x56t6r9vh9xjmg5whtkmp'
    # Tone ID: px4jkk2nqrf9h6gp7wwxnhvz
    # Location ID: nm6xqgcf5j7sw8w994c6sc8h
    alt_tag_id = 'vk54n6pwnxm27gwrvrzfb'
    topic_id = '9mcg3rgm5mf72jqrtw2gqm7t'
    project_name_id = '5zpqwt2r348sjbnc6rpxc96'
    caption_id = 'cmcbhcc5nmm72v57vrxppw2x'
    # Original Project Images Section ID: c5vm8cnh9jvkjbh7r43qxkv
    # Edited Project Images Section ID: 5wpz2s9m3g7ctcjpm4vrt46

    r_asset = requests.post(f'https://brandfolder.com/api/v4/assets/{asset_id}/tags', json=tags_payload, params={}, headers=headers)

    # alt_tags
    r_topic = requests.post(f'https://brandfolder.com/api/v4/custom_field_keys/{topic_id}/custom_field_values', json=
          topic_payload
        , params={
      }, headers=headers)

    r_alt_tag = requests.post(f'https://brandfolder.com/api/v4/custom_field_keys/{alt_tag_id}/custom_field_values', json=
          alt_tag_payload
        , params={
      }, headers=headers)

    r_year = requests.post(f'https://brandfolder.com/api/v4/custom_field_keys/{year_id}/custom_field_values', json=
          year_payload
        , params={
      }, headers=headers)

    r_client = requests.post(f'https://brandfolder.com/api/v4/custom_field_keys/{client_name_id}/custom_field_values', json=
          client_payload
        , params={
      }, headers=headers)

    r_project = requests.post(f'https://brandfolder.com/api/v4/custom_field_keys/{project_name_id}/custom_field_values', json=
          project_payload
        , params={
      }, headers=headers)

    r_caption = requests.post(f'https://brandfolder.com/api/v4/custom_field_keys/{caption_id}/custom_field_values', json=
        caption_payload
      , params={
    }, headers=headers)

    return image_name

def delete_og_asset(asset_id):
  headers = {
      'Accept': 'application/json',
      'Authorization': 'eyJhbGciOiJIUzI1NiJ9.eyJvcmdhbml6YXRpb25fa2V5IjoiZmY0cmt0NDNoMzRtMjVoa2duNWJteDlmIiwiaWF0IjoxNzA1OTQ4NjI3LCJ1c2VyX2tleSI6IjhyNnhxeDR6bTdyN2Z4NnJqY25jM2IzIiwic3VwZXJ1c2VyIjpmYWxzZX0.xUPT9j08a0THBwW_0GkQjllJxmjeDGtcPeoIOu_w9Zs'
    }

  r = requests.delete(f'https://brandfolder.com/api/v4/assets/{asset_id}', params={
      # use a dict with your desired URL parameters here
  }, headers=headers)

  return

def run_preprocess_ai(topical_map, client_name, section_type, query_engine=None, progress=gr.Progress()):
  section_id = section_key_dict[section_type]
  headers = {
    'Content-Type': 'application/json',
    'Authorization': brandfolder_api
  }
  collection_dict = get_collection_dict()
  collection_id = collection_dict[client_name]
  page = 1
  pre_process_ids = []
  run = True
  while run == True:
    r = requests.get(f'https://brandfolder.com/api/v4/collections/{collection_id}/assets?include=section,custom_fields&fields=created_at&page={page}&per=3000&sort_by=created_at&order=DESC', params={}, headers=headers)    
    page+=1
    asset_names = [item['id'] for item in r.json()['data'] if item['relationships']['section']['data']['id'] == section_id]
    if asset_names in pre_process_ids:
      run = False
    else:
        pre_process_ids.append(asset_names)
  asset_names = sum(pre_process_ids, [])
  
  new_imgs = []
  for asset_id in progress.tqdm(asset_names, desc="Uploading..."):
    try:
        time.sleep(2)
        asset_dict = get_asset_info(asset_id)
        new_img = create_ai_asset(asset_dict, topical_map, client_name, new_imgs, query_engine=query_engine[-1])
        new_imgs.append(new_img)
        if new_img:
          delete_og_asset(asset_id)
    except Exception as e:
        print(f'An unexpected error occured processing {asset_dict["asset_name"]}: {e}')
  gr.Info('Images have been processed!')
  return 'Images Processed'