import base64
import time

import requests
import cv2
import json
import re
import os
import tqdm

# OpenAI API Key
api_key = "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"

# Function to encode the image
def encode_image(image_path):
  img = cv2.imread(image_path)
  scale = 500.0 / min(img.shape[:2])
  img = cv2.resize(img, (0, 0), fx=scale, fy=scale)
  scaled_path = '/tmp/scaled_image.jpg'
  cv2.imwrite(scaled_path, img)
  with open(scaled_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')


def encode_image_v2(image_in):
  img = image_in
  scale = 500.0 / min(img.shape[:2])
  img = cv2.resize(img, (0, 0), fx=scale, fy=scale)
  scaled_path = '/tmp/scaled_image.jpg'
  cv2.imwrite(scaled_path, img)
  with open(scaled_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')

def caption_image(image_in):
  # Getting the base64 string
  base64_image = encode_image_v2(image_in)

  headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {api_key}"
  }

  payload = {
    "model": "gpt-4-vision-preview",
    "messages": [
      {
        "role": "user",
        "content": [
          {
            "type": "text",
            "text": "As an AI image tagging expert, please provide precise tags for the hairstyle in the image, To enhance CLIP model's understanding of the content. Please provide a detailed description of the hairstyle in the image, including but not limited to the color, style, length, curliness, hairline, highlights, gradients,  etc. Your tags should be accurate, non-duplicative, and within a 10-20 word count range. Tags should be comma-separated. No need to provide any safety statements or precautions."
          },
          {
            "type": "image_url",
            "image_url": {
              "url": f"data:image/jpeg;base64,{base64_image}",
              "detail": "low"
            }
          }
        ]
      }
    ],
    "max_tokens": 300
  }

  # try:
  #   response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
  #
  #   tmp_json = response.json()
  #   return tmp_json['choices'][0]['message']['content']
  # except Exception as e:
  #   print(e)
  #   return ""
  return ""



if __name__ == '__main__':
  prompt = caption_image('/home/chinatszrn/Downloads/abc/train_data/style1/07ebac82-4c0f-4dd2-84bd-bc34a059bd9b.png')
  print(prompt)