import os
import json
import gradio as gr
import requests
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()
API_KEY = os.getenv("DEEPSEEK_API_KEY")
API_KEY = "sk-0d040569ee2b4932a807d56a5b12ac05"
API_URL = "https://api.deepseek.com/v1/chat/completions"

def build_messages(query: str, history: list):
    """构造对话历史"""
    messages = [{"role": "system", "content": "你是一个有帮助的助手"}]
    for user_msg, bot_msg in history:
        messages.extend([
            {"role": "user", "content": user_msg},
            {"role": "assistant", "content": bot_msg}
        ])
    messages.append({"role": "user", "content": query})
    return messages

def stream_response(query: str, history: list):
    """流式响应生成器"""
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            API_URL,
            headers=headers,
            json={
                "model": "deepseek-chat",
                "messages": build_messages(query, history),
                "temperature": 0.7,
                "stream": True
            },
            stream=True,
            timeout=60
        )
        
        # 先检查响应状态再处理内容
        if response.status_code != 200:
            error_msg = f"API 返回错误状态码: {response.status_code} - {response.text}"
            print(error_msg)
            yield error_msg
            return

        partial_message = ""
        for chunk in response.iter_lines():
            # 过滤空行和 keep-alive 消息
            if not chunk or b'[DONE]' in chunk:
                continue

            try:
                # 调试输出原始数据
                print("原始 chunk:", chunk)
                
                decoded = chunk.decode('utf-8').strip()
                # 处理可能的多个 data: 前缀
                if decoded.startswith('data:'):
                    json_str = decoded[5:].strip()
                else:
                    json_str = decoded
                
                # 验证 JSON 有效性
                if not json_str.startswith('{'):
                    continue
                
                data = json.loads(json_str)
                
                if content := data['choices'][0]['delta'].get('content'):
                    partial_message += content
                    yield partial_message

            except json.JSONDecodeError as e:
                print(f"JSON 解析失败: {e} | 原始数据: {decoded}")
                continue
            except Exception as e:
                print(f"处理 chunk 时发生错误: {str(e)}")
                continue

    except Exception as e:
        yield f"⚠️ 请求失败：{str(e)}"

def stream_response_back(query: str, history: list):
    """流式响应生成器"""
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    print(query)
    print(history)

    try:
        response = requests.post(
            API_URL,
            headers=headers,
            json={
                "model": "deepseek-chat",
                "messages": build_messages(query, history),
                "temperature": 0.7,
                "stream": True
            },
            stream=True,
            timeout=60
        )
        response.raise_for_status()
        
        print(response)

        partial_message = ""
        for chunk in response.iter_lines():
            if chunk:
                # 处理流式数据格式
                decoded = chunk.decode('utf-8').strip()
                if decoded.startswith('data:'):
                    data = json.loads(decoded[5:])
                    if content := data['choices'][0]['delta'].get('content'):
                        partial_message += content
                        yield partial_message

    except Exception as e:
        yield f"⚠️ 请求失败：{str(e)}"

# 创建带打字机效果的聊天界面
demo = gr.ChatInterface(
    fn=stream_response,
    title="DeepSeek 智能助手",
    description="输入消息开始对话（支持流式打字效果）",
    theme="soft",
    examples=["你好！", "如何学习AI？", "写一首关于春天的诗"],
    cache_examples=False,
    # retry_btn=None,
    # undo_btn=None,
    # clear_btn="清空历史",
    stop_btn="停止生成",
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)