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
from pillow_heif import register_heif_opener
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")

  register_heif_opener()
  
  def compress_and_encode_image(url, target_size_mb=20, quality=70):
      # Fetch the image from the URL
      # Open the image using Pillow
      img = Image.open(requests.get(url, stream=True).raw)
      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)

# creates the asset in the client's brand folder
def create_ai_asset(asset_dict, topical_map, collection_name, new_imgs, 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']

    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']

    url_response = urllib.request.urlopen(image_url)
    img_array = np.array(bytearray(url_response.read()), dtype=np.uint8)
    image = cv2.imdecode(img_array, -1)
    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))

    # image = sharpen_image(image)

    with NamedTemporaryFile(delete=True, suffix='.jpg') as temp_image:
      # fp = TemporaryFile()
      cv2.imwrite(temp_image.name, image)
      # fp.seek(0)
      response = requests.put(upload_url, data=temp_image)
    # fp.close()

    # 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, 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:
    if section_type == 'Pre-Processed Images':
        r = requests.get(f'https://brandfolder.com/api/v4/brandfolders/988cgqcg8xsrr5g9h7gtsqkg/assets?include=section,custom_fields&fields=created_at&page={page}&per=3000&sort_by=created_at&order=DESC', params={}, headers=headers)
    elif section_type == 'Original Project Assets':
        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)
    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