import os
import json
import asyncio
import httpx
import traceback
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

AIROPS_API_KEY = os.environ.get("AIROPS_API_KEY", "")
AIROPS_APP_ID = os.environ.get("AIROPS_APP_ID", "")
INPUT_VAR = os.environ.get("INPUT_VAR_NAME", "input")

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        try:
            body = await request.json()
        except:
            body = {}
            
        messages = body.get("messages", [])
        
        # 将所有历史消息和设定拼接成一个长剧本
        full_context = ""
        for msg in messages:
            # 强制转换为字符串，防止特殊格式报错
            role = str(msg.get("role", "user")).upper()
            content = str(msg.get("content", ""))
            
            if role == "SYSTEM":
                full_context += f"[系统设定及角色卡]:\n{content}\n\n"
            elif role == "ASSISTANT":
                full_context += f"[角色回复]:\n{content}\n\n"
            else:
                full_context += f"[我的输入]:\n{content}\n\n"
                
        if not full_context.strip():
            full_context = "[我的输入]:\n你好"
        
        headers = {
            "Authorization": f"Bearer {AIROPS_API_KEY}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "inputs": {
                INPUT_VAR: full_context.strip()
            }
        }
        
        async def generate():
            # 同样保持 300 秒超时，给 AI 充分的思考时间
            async with httpx.AsyncClient(timeout=300.0) as client:
                url = f"https://app.airops.com/public_api/airops_apps/{AIROPS_APP_ID}/execute"
                
                try:
                    response = await client.post(url, headers=headers, json=payload)
                    
                    if response.status_code == 200:
                        data = response.json()
                        result_text = data.get("result", "")
                        
                        if isinstance(result_text, (dict, list)):
                            result_text = json.dumps(result_text, ensure_ascii=False)
                        
                        # 模拟 OpenAI 流式打字机效果
                        chunk_size = 3
                        for i in range(0, len(result_text), chunk_size):
                            chunk = result_text[i:i+chunk_size]
                            openai_chunk = {
                                "id": "chatcmpl-airops",
                                "object": "chat.completion.chunk",
                                "choices": [{"delta": {"content": chunk}}]
                            }
                            yield f"data: {json.dumps(openai_chunk)}\n\n"
                            await asyncio.sleep(0.01)
                            
                        yield "data: [DONE]\n\n"
                    else:
                        print(f"AirOps Error: {response.status_code} - {response.text}")
                        err = {"error": f"AirOps 接口报错: {response.status_code}"}
                        yield f"data: {json.dumps(err)}\n\n"
                        yield "data: [DONE]\n\n"
                except Exception as e:
                    traceback.print_exc()
                    err = {"error": f"代理服务器内部错误: {repr(e)}"}
                    yield f"data: {json.dumps(err)}\n\n"
                    yield "data: [DONE]\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")
        
    except Exception as e:
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"error": f"请求解析错误: {repr(e)}"})