import base64
#from openai import OpenAI

# client = OpenAI(
#     #api_key="sk-6Fr1OWm8SJXseEtWjArObksjTr39P08HSNydgNFgsOewI0xp",
#     api_key="sk-zLFhpbBQKWV4qohV4vD0AFR6Gaa7JwL0Q8zphh1CYFl2EqDM",
#     base_url="https://www.chataiapi.com/v1",
# )
#https://ai.juguang.chat/console/topup
import requests
import json
import http.client
import json
import base64
import re

def save_base64_image_from_json(json_str, output_filename):
    """
    从JSON字符串中提取base64图片数据并保存为文件
    """
    try:
        # 解析JSON
        data = json.loads(json_str)
        
        parts = data["candidates"][0]["content"]["parts"]
        
        # 遍历parts找到包含inlineData的那个
        base64_data = None
        for part in parts:
            if "inlineData" in part and "data" in part["inlineData"]:
                base64_data = part["inlineData"]["data"]
                break
        
        if base64_data is None:
            print("未找到包含图片数据的部分")
            return False
        
        # 清理可能的换行符和空格
        base64_data = re.sub(r'\s+', '', base64_data)
        
        # 解码base64数据
        image_data = base64.b64decode(base64_data)
        
        # 保存为文件
        with open(output_filename, 'wb') as f:
            f.write(image_data)
        
        print(f"图片已成功保存为: {output_filename}")
        return True
        
    except Exception as e:
        print(f"保存图片时发生错误: {e}")
        return False

def image_to_base64(file_path, mime_type=None):
    if mime_type is None:
        extension = file_path.split('.')[-1].lower()
        mime_types = {
            'jpg': 'image/jpeg',
            'jpeg': 'image/jpeg',
            'png': 'image/png',
            'gif': 'image/gif',
            'webp': 'image/webp',
            'bmp': 'image/bmp'
        }
        mime_type = mime_types.get(extension, 'application/octet-stream')
    
    # 读取文件内容并编码为Base64
    with open(file_path, 'rb') as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    
    # 组合成Data URI格式
    #return f"data:{mime_type};base64,{encoded_string}"
    return encoded_string 

conn = http.client.HTTPSConnection("ai.juguang.chat")
payload = json.dumps({
  "contents": [
    {
      "parts": [
        {
          "text": "这是一张有人的图片"
        },
        {
          "inline_data": {
            "mime_type": "image/jpeg",
            "data": image_to_base64("/home/xsl/code/ComfyUI/change_cloth/static/imgs/0b9014f9-14dc-4a97-9c99-642441b8c122.jpg")
          }
        },
        {
          "text": "这是另外一张有衣服，裤子或者裙子，和其他各种可以穿戴在身上物品的图片"
        },
        {
          "inline_data": {
            "mime_type": "image/jpeg",
            "data": image_to_base64("/home/xsl/code/ComfyUI/change_cloth/cloth_suit.jpg")
          }
        },
        {
          "text": "给第一张图片的人物，换上第二张图片的衣服和配饰"
        }
      ]
    }
  ]
})
headers = {
  'Content-Type': 'application/json',
  'Authorization': 'Bearer sk-zLFhpbBQKWV4qohV4vD0AFR6Gaa7JwL0Q8zphh1CYFl2EqDM'
}
conn.request("POST", "/v1beta/models/gemini-2.5-flash-image-preview:generateContent", payload, headers)
res = conn.getresponse()
data = res.read()
decode_str = data.decode("utf-8")


# 调用函数保存图片
save_base64_image_from_json(decode_str, "out_image.png")

print(decode_str)