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):
    # 验证 GPU 是否可用
    print(f"当前设备: {zero.device}", file=sys.stderr)
    
    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***",  # 可以根据需要修改用户ID
        "stream": True,
        "auto_save_history": True,
        "additional_messages": messages
    }
    
    print(f"发送到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)
            if isinstance(chunk, dict) and "content" in chunk and chunk.get("type") == "answer":
                full_response = chunk["content"]
        except json.JSONDecodeError:
            pass
    
    print(f"完整的响应: {full_response}", file=sys.stderr)
    return full_response

iface = gr.ChatInterface(
    chat_stream,
    chatbot=gr.Chatbot(height=600, label="提示词助手"),
    textbox=gr.Textbox(placeholder="在这里输入您的消息...", container=False, scale=7),
    title="米你AI 提示词助手",
    description="帮助你生成提示词，输入一个主题，我会帮你联想合适的提示词",
    theme="soft",
    examples=[
        "冬至主题的海报",
        "温馨的插画，男孩和宠物主题",
        "一个穿着太空服的小鸡，在月球上行走"
    ],
    cache_examples=False,
    retry_btn="重试",
    undo_btn="撤销",
    clear_btn="清除",
    submit_btn="发送",
    stop_btn="停止",
    autofocus=True
)

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