import gradio as gr
import requests
import os
from google.oauth2 import service_account
from google.cloud import storage




# 读取图片
import base64
# Function to encode the image
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')


openai_api_key = os.environ.get('openai_api_key')

# 将图片上传到google cloud storage
def upload_image_to_gcs_blob(image):

    google_creds = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")

    creds_json = json.loads(google_creds)
    credentials = service_account.Credentials.from_service_account_info(creds_json)

    # 现在您可以使用这些凭证对Google Cloud服务进行认证
    storage_client = storage.Client(credentials=credentials, project=creds_json['project_id'])

    bucket_name = os.environ.get('bucket_name')
    bucket = storage_client.bucket(bucket_name)
    
    destination_blob_name = os.path.basename(image)
    blob = bucket.blob(destination_blob_name)

    blob.upload_from_filename(image)

    public_url = blob.public_url
    
    return public_url


def ask_image(text,image,api_token=openai_api_key):
    public_url = upload_image_to_gcs_blob(image)
    messages=[
      {
        "role": "user",
        "content": [
          {"type": "text", "text": text},
          {
            "type": "image_url",
            "image_url": {
                # "url":f"data:image/jpeg;base64,{base64_image}"
                "url": public_url
            },
          },
        ],
      }
    ]
    
    # 请求头部信息
    headers = {
      'Authorization': f'Bearer {api_token}'
    }
    
    # 请求体信息
    data = {
      'model': 'gpt-4o',  # 可以根据需要更换其他模型
      'messages': messages,
      'temperature': 0.7  # 可以根据需要调整
    }
    
    
    # 设定最大重试次数
    max_retry = 3
    
    for i in range(max_retry):
      try:
          # 发送请求
          response = requests.post('https://burn.hair/v1/chat/completions', headers=headers, json=data)
    
          # 解析响应内容
          response_data = response.json()
          response_content = response_data['choices'][0]['message']['content']
          usage = response_data['usage']
    
          # response_content = 'test response'
    
          return response_content
    
      except Exception as e:
          # 如果已经达到最大重试次数，那么返回空值
          if i == max_retry - 1:
              print(f'重试次数已达上限，仍未能成功获取数据，错误信息：{e}')
              response_content = ''
              usage = {}
              return response_content
          else:
              # 如果未达到最大重试次数，打印错误信息，并继续下一次循环
              print(f'第{i+1}次请求失败，错误信息：{e}，准备进行第{i+2}次尝试')


# gradio demo

title = "Ask Image"
description = "Ask anything about your Image"

demo = gr.Interface(
    fn=ask_image,
    inputs=[gr.Text(label="Question"),gr.Image(label='',type='filepath')],
    outputs=[gr.Textbox(label="Answer",lines=3)],
    title = title,
    description = description
)
demo.queue(max_size = 20)

demo.launch(share = True)