import os
import json
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from openai import OpenAI
from dotenv import load_dotenv
from starlette.middleware.cors import CORSMiddleware

load_dotenv()

app = FastAPI()

client = OpenAI(
    api_key=os.environ.get("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],   # 生产环境改成具体域名
    allow_methods=["*"],
    allow_headers=["*"],
)


def stream_generator(user_input: str, model: str):
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "user", "content": user_input},
            ],
            stream=True,
            # reasoning_effort 与 thinking 参数需按当前 DeepSeek 官方文档核验
            reasoning_effort="high",
            extra_body={"thinking": {"type": "enabled"}},
        )

        for chunk in response:
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            import time
            time.sleep(0.5)

            reasoning = getattr(delta, "reasoning_content", None)
            if reasoning:
                yield f"event: reasoning\ndata: {json.dumps(reasoning, ensure_ascii=False)}\n\n"

            content = getattr(delta, "content", None)
            if content:
                yield f"event: content\ndata: {json.dumps(content, ensure_ascii=False)}\n\n"

        yield "event: done\ndata: [DONE]\n\n"

    except Exception as e:
        yield f"event: server_error\ndata: {json.dumps(str(e), ensure_ascii=False)}\n\n"


@app.get("/chat")
def call_llm(user_input: str, model: str = "deepseek-chat"):
    return StreamingResponse(
        stream_generator(user_input, model),
        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=8001)