import gradio as gr
import requests
import json
import sseclient
import sys
import os

COZE_API_URL = "https://api.coze.cn/v3/chat"
API_KEY = os.environ.get("COZE_API_KEY", "your_default_key")
BOT_ID = os.environ.get("COZE_BOT_ID", "your_default_bot_id")

def chat_stream(message, history):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    # 构建包含历史记录的消息列表
    messages = []
    for human, assistant in history:
        messages.append({"role": "user", "content": human, "content_type": "text"})
        messages.append({"role": "assistant", "content": assistant, "content_type": "text"})
    messages.append({"role": "user", "content": message, "content_type": "text"})
    
    data = {
        "bot_id": BOT_ID,
        "user_id": "123123***",
        "stream": True,
        "auto_save_history": True,
        "additional_messages": messages
    }
    
    print(f"\n[请求] 发送到API的数据: {json.dumps(data, ensure_ascii=False)}", file=sys.stderr)
    
    response = requests.post(COZE_API_URL, headers=headers, data=json.dumps(data), stream=True)
    client = sseclient.SSEClient(response)
    
    full_response = ""
    for event in client.events():
        if event.data == "[DONE]":
            break
        try:
            chunk = json.loads(event.data)
            print(f"\n[调试] 收到的事件数据: {event.data[:200]}...", file=sys.stderr)
            
            if isinstance(chunk, dict):
                if chunk.get("type") == "answer":
                    content = json.loads(chunk["content"])
                    print("\n[处理] 解析 content 内容", file=sys.stderr)
                    
                    if content.get("card_type") == 3:
                        data_str = content.get("data", "")
                        try:
                            data_json = json.loads(data_str)
                            if "variables" in data_json:
                                for variable in data_json["variables"].values():
                                    if variable.get("name") == "img_url":
                                        img_url = variable.get("defaultValue")
                                        if img_url:
                                            full_response = f'''
                                            <img src="{img_url}" 
                                                 width="600px" 
                                                 onclick="window.open(this.src, '_blank')" 
                                                 style="cursor: pointer;" 
                                                 title="点击查看大图">
                                            '''
                                            print(f"\n[成功] 提取到图片URL: {img_url}", file=sys.stderr)
                                            print(f"\n[成功] 构建的图片标签: {full_response}", file=sys.stderr)
                                            break
                        except json.JSONDecodeError as e:
                            print(f"\n[错误] 解析 data 字段失败: {e}", file=sys.stderr)
        except json.JSONDecodeError as e:
            print(f"\n[错误] 解析事件数据失败: {e}", file=sys.stderr)
    
    print(f"\n[输出] 最终返回到对话框的响应: {full_response}", file=sys.stderr)
    return full_response

iface = gr.ChatInterface(
    chat_stream,
    css="""
    .message-wrap img { 
        max-width: 400px !important; 
        height: auto !important;
        max-height: unset !important;
        transition: transform 0.2s;
    }
    .message-wrap img:hover {
        transform: scale(1.02);
    }
    """,
    chatbot=gr.Chatbot(height=800, label="提示词助手"),
    textbox=gr.Textbox(placeholder="在这里输入您的消息...", container=False, scale=7),
    title="米你AI 海报助手",
    description="告诉我你的海报需求，我会帮你提取关键信息、生成海报",
    theme="soft",
    examples=[
        "苹果滞销！老农都急哭了！现在 5 折甩卖。写实风格",
        "中秋节到了，团建活动来啦！我们提供了丰富的食物和游戏",
        "儿童滑板大酬宾啦，现在下单，优惠多多！"
    ],
    cache_examples=False,
    retry_btn="重试",
    undo_btn="撤销",
    clear_btn="清除",
    submit_btn="发送",
    stop_btn="停止",
    autofocus=True
)

if __name__ == "__main__":
    iface.launch(share=True)