import os
import json
import asyncio
import httpx
from typing import AsyncGenerator, Optional
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

# 初始化 FastAPI 应用
app = FastAPI(title="Qwen3.6 Thinking Stream Proxy")

# 配置 CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Qwen3.6 服务配置
# 假设使用兼容 OpenAI 格式的接口，如本地部署或 DashScope
QWEN_BASE_URL = os.getenv("QWEN_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
QWEN_API_KEY = os.getenv("QWEN_API_KEY", "your_api_key_here")
QWEN_MODEL = os.getenv("QWEN_MODEL", "qwen-plus") # 或 qwen-max, qwen-turbo 等支持思考的模型

class ChatRequest(BaseModel):
    message: str
    model: Optional[str] = QWEN_MODEL
    enable_thinking: bool = True # 默认开启思考模式

async def event_generator(request_body: dict, client: httpx.AsyncClient) -> AsyncGenerator[str, None]:
    """
    异步生成器：从 Qwen 接口获取 SSE 流并转发给前端
    """
    headers = {
        "Authorization": f"Bearer {QWEN_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream"
    }
    
    # 构建符合 OpenAI 兼容格式的请求体
    # 注意：enable_thinking 通常放在 extra_body 或特定参数中，具体取决于服务商实现
    # 这里假设通过 extra_body 传递，或者如果是 DashScope 原生接口需调整结构
    # 以下以 OpenAI 兼容格式为例，若使用 DashScope 原生 SDK 需调整 payload 结构
    payload = {
        "model": request_body.get("model", QWEN_MODEL),
        "messages": [
            {"role": "user", "content": request_body["message"]}
        ],
        "stream": True,
        "stream_options": {"include_usage": False} # 可选
    }
    
    # 如果使用的是 DashScope 兼容模式且支持 extra_body 传递 thinking 参数
    if request_body.get("enable_thinking", True):
        payload["extra_body"] = {"enable_thinking": True}

    try:
        async with client.stream(
            "POST", 
            f"{QWEN_BASE_URL}/chat/completions", 
            json=payload, 
            headers=headers,
            timeout=None
        ) as response:
            if response.status_code != 200:
                error_detail = await response.aread()
                raise HTTPException(status_code=response.status_code, detail=f"Upstream error: {error_detail.decode()}")
            
            async for line in response.aiter_lines():
                if not line:
                    continue
                
                # 处理 SSE 结束标记
                if line.strip() == "[DONE]":
                    yield "data: [DONE]\n\n"
                    break
                
                # 解析上游返回的 JSON 数据
                if line.startswith("data: "):
                    json_str = line[6:]
                    try:
                        data = json.loads(json_str)
                        
                        # 提取内容片段
                        # Qwen 的思考过程通常在 delta.content 或 delta.reasoning_content 中
                        # 具体字段名取决于模型版本和接口定义
                        choices = data.get("choices", [])
                        if choices:
                            delta = choices.get("delta", {})
                            
                            # 尝试获取思考内容 (reasoning_content) 和普通内容 (content)
                            reasoning = delta.get("reasoning_content", "")
                            content = delta.get("content", "")
                            
                            # 构造前端易于区分的数据格式
                            # 例如：{ "type": "thinking", "data": "..." } 或 { "type": "answer", "data": "..." }
                            if reasoning:
                                output_data = json.dumps({"type": "thinking", "content": reasoning}, ensure_ascii=False)
                                yield f"data: {output_data}\n\n"
                            
                            if content:
                                output_data = json.dumps({"type": "answer", "content": content}, ensure_ascii=False)
                                yield f"data: {output_data}\n\n"
                                
                    except json.JSONDecodeError:
                        # 如果不是 JSON，直接透传（以防万一）
                        yield f"data: {json_str}\n\n"
                else:
                    # 非 data: 开头的行，可能是注释或其他，忽略或按需处理
                    pass

    except httpx.RequestError as exc:
        yield f"data: {{\"type\": \"error\", \"content\": \"Connection error: {str(exc)}\"}}\n\n"
    except Exception as exc:
        yield f"data: {{\"type\": \"error\", \"content\": \"Internal error: {str(exc)}\"}}\n\n"

@app.post("/api/chat/stream")
async def chat_stream(request: ChatRequest):
    """
    代理接口：接收前端提问，转发给 Qwen3.6，流式返回思考和回答
    """
    upstream_payload = request.dict()
    
    async with httpx.AsyncClient() as client:
        return StreamingResponse(
            event_generator(upstream_payload, client),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"
            }
        )

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