"""
FastAPI SSE 接口 - Qwen 模型思维链展示
实时返回模型的思考过程（think process）
"""
import asyncio
import json
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import httpx

app = FastAPI(
    title="Qwen Think Process API",
    description="展示 Qwen 模型的思维链（Think Process），实时返回思考过程",
    version="1.0.0"
)

# Qwen API 配置（预留配置参数）
QWEN_API_URL = "http://localhost:8000/v1/chat/completions"  # Qwen 服务地址
QWEN_API_KEY = "your-api-key-here"  # API Key 预留
QWEN_MODEL = "Qwen/Qwen2-7B-Instruct"  # 模型名称


class ThinkProcessRequest(BaseModel):
    """请求体模型"""
    question: str
    max_tokens: int = 500
    model: str = QWEN_MODEL
    api_key: str = None  # 允许通过请求覆盖 API Key


async def call_qwen_api(question: str, model: str = None, api_key: str = None):
    """
    真实调用 Qwen API 获取思维链
    
    Args:
        question: 用户问题
        model: 模型名称（可选，默认使用配置值）
        api_key: API Key（可选，默认使用配置值）
    """
    # 使用传入值或默认配置
    target_model = model or QWEN_MODEL
    target_api_key = api_key or QWEN_API_KEY
    
    # 构建提示词，引导模型输出思维链
    think_prompt = f"""
    请用中文详细展示你的思考过程来回答以下问题：
    
    问题：{question}
    
    思考过程格式：
    1. 分析问题核心需求
    2. 检索相关知识
    3. 构建推理链条
    4. 验证逻辑正确性
    5. 生成最终回答
    
    请按步骤详细说明你的思考过程。
    """
    
    payload = {
        "model": target_model,
        "messages": [
            {
                "role": "system",
                "content": "你是一个擅长展示思考过程的AI助手，请详细展示你的思考步骤。"
            },
            {
                "role": "user",
                "content": think_prompt
            }
        ],
        "stream": True,
        "max_tokens": 1000,
        "temperature": 0.7
    }
    
    headers = {
        "Content-Type": "application/json"
    }
    
    # 如果有 API Key，添加认证头
    if target_api_key and target_api_key != "your-api-key-here":
        headers["Authorization"] = f"Bearer {target_api_key}"
    
    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            async with client.stream(
                "POST",
                QWEN_API_URL,
                json=payload,
                headers=headers
            ) as response:
                if response.status_code != 200:
                    try:
                        error_data = await response.json()
                        error_msg = error_data.get("error", {}).get("message", f"API返回错误: {response.status_code}")
                    except Exception:
                        error_msg = f"API请求失败: {response.status_code}"
                    
                    yield {"type": "error", "content": error_msg}
                    yield {"type": "completed", "content": ""}
                    return
                
                # 解析流式响应
                thinking_content = ""
                async for chunk in response.aiter_bytes(chunk_size=64):
                    try:
                        chunk_text = chunk.decode("utf-8", errors="ignore")
                        # 解析 SSE 格式
                        for line in chunk_text.split("\n"):
                            if line.startswith("data: "):
                                data_str = line[6:].strip()
                                if data_str == "[DONE]":
                                    continue
                                try:
                                    data = json.loads(data_str)
                                    choices = data.get("choices", [])
                                    if choices:
                                        delta = choices[0].get("delta", {})
                                        content = delta.get("content", "")
                                        if content:
                                            thinking_content += content
                                            yield {
                                                "type": "thinking",
                                                "content": content
                                            }
                                except json.JSONDecodeError:
                                    continue
                    except Exception as e:
                        continue
                
                # 最终回答汇总
                if thinking_content:
                    yield {"type": "answer", "content": thinking_content}
                
                yield {"type": "completed", "content": ""}
    
    except httpx.TimeoutException:
        yield {"type": "error", "content": "API请求超时"}
        yield {"type": "completed", "content": ""}
    except httpx.ConnectError:
        yield {"type": "error", "content": "无法连接到Qwen服务"}
        yield {"type": "completed", "content": ""}
    except Exception as e:
        yield {"type": "error", "content": f"API调用异常: {str(e)}"}
        yield {"type": "completed", "content": ""}


@app.post("/api/v1/think")
async def think_stream(request: Request):
    """
    SSE 接口：实时返回 Qwen 模型的思维链
    
    请求体：
    {
        "question": "您的问题",
        "max_tokens": 500,
        "model": "Qwen/Qwen2-7B-Instruct",  # 可选，覆盖默认模型
        "api_key": "your-key"  # 可选，覆盖默认API Key
    }
    
    响应格式（SSE）：
    - thinking: 思考过程（流式返回）
    - answer: 最终回答汇总
    - error: 错误信息
    - completed: 完成标记
    """
    # 解析请求体
    body = await request.body()
    try:
        data = json.loads(body.decode("utf-8"))
        question = data.get("question", "")
        model = data.get("model", None)
        api_key = data.get("api_key", None)
    except Exception:
        question = "用户提问"
        model = None
        api_key = None
    
    async def stream_generator():
        async for step in call_qwen_api(question, model, api_key):
            yield f"data: {json.dumps(step, ensure_ascii=False)}\n\n"
    
    return StreamingResponse(
        stream_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
        }
    )


@app.get("/")
async def root():
    return {
        "service": "Qwen Think Process API",
        "version": "1.0.0",
        "endpoints": {
            "think": "POST /api/v1/think (SSE)"
        },
        "description": "实时展示 Qwen 模型的思维链过程",
        "config": {
            "api_url": QWEN_API_URL,
            "default_model": QWEN_MODEL
        }
    }


@app.get("/health")
async def health():
    return {"status": "healthy"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("qwen:app", host="0.0.0.0", port=8000, reload=True)