# requires: fastapi, httpx
"""流式对话 API：鉴权、参数校验、请求 ID、统一错误格式、SSE 流式输出、依赖注入替换模型客户端。

运行：uvicorn app:app --reload（或 python app.py 打印一次演示）
测试：python -m pytest -q test_app.py
模型客户端是假的（逐字返回固定文本），换成真实 SDK 时只需要实现同样的 stream 方法。
"""
from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from collections.abc import AsyncIterator
from typing import Annotated, Literal, Protocol

from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

log = logging.getLogger("chat-api")


# ---------- 模型 ----------
class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str = Field(min_length=1, max_length=4000)


class ChatRequest(BaseModel):
    messages: list[Message] = Field(min_length=1, max_length=50)
    temperature: float = Field(default=0.3, ge=0, le=2)
    stream: bool = True


class ChatResponse(BaseModel):
    request_id: str
    content: str


# ---------- 依赖 ----------
class LLMClient(Protocol):
    def stream(self, messages: list[Message], temperature: float) -> AsyncIterator[str]: ...


class FakeLLM:
    async def stream(self, messages: list[Message], temperature: float) -> AsyncIterator[str]:
        answer = f"你问的是：{messages[-1].content}。这是一个演示回答。"
        for ch in answer:
            await asyncio.sleep(0.005)
            yield ch


def get_llm() -> LLMClient:
    return FakeLLM()


API_KEYS = {"demo-key": "tenant-a"}  # 演示用；真实项目从数据库或密钥服务读取，并做哈希比对


def current_tenant(x_api_key: Annotated[str | None, Header()] = None) -> str:
    tenant = API_KEYS.get(x_api_key or "")
    if tenant is None:
        raise HTTPException(status_code=401, detail="API Key 无效")
    return tenant


# ---------- 应用 ----------
app = FastAPI(title="chat-api")


@app.middleware("http")
async def request_id_and_timing(request: Request, call_next):
    rid = request.headers.get("X-Request-ID") or uuid.uuid4().hex[:16]
    request.state.request_id = rid
    start = time.perf_counter()
    response = await call_next(request)
    response.headers["X-Request-ID"] = rid
    log.info("%s %s %s %.1fms", request.method, request.url.path, response.status_code,
             (time.perf_counter() - start) * 1000)
    return response


def error_body(request: Request, code: str, message: str, details=None) -> dict:
    body = {"code": code, "message": message, "request_id": getattr(request.state, "request_id", "")}
    if details is not None:
        body["details"] = details
    return body


@app.exception_handler(HTTPException)
async def http_error(request: Request, exc: HTTPException):
    return JSONResponse(error_body(request, f"HTTP_{exc.status_code}", str(exc.detail)), status_code=exc.status_code)


@app.exception_handler(RequestValidationError)
async def validation_error(request: Request, exc: RequestValidationError):
    details = [{"field": ".".join(str(p) for p in e["loc"][1:]), "message": e["msg"]} for e in exc.errors()]
    return JSONResponse(error_body(request, "INVALID_PARAMS", "参数校验失败", details), status_code=422)


@app.exception_handler(Exception)
async def unhandled_error(request: Request, exc: Exception):
    log.exception("未处理的异常")  # 详细信息只进日志，不返回给调用方
    return JSONResponse(error_body(request, "INTERNAL", "服务器内部错误"), status_code=500)


@app.get("/healthz")
async def healthz() -> dict:
    return {"status": "ok"}


@app.post("/v1/chat", response_model=None)
async def chat(req: ChatRequest, request: Request,
               tenant: Annotated[str, Depends(current_tenant)],
               llm: Annotated[LLMClient, Depends(get_llm)]):
    rid = request.state.request_id
    if not req.stream:
        content = "".join([t async for t in llm.stream(req.messages, req.temperature)])
        return ChatResponse(request_id=rid, content=content)

    async def events() -> AsyncIterator[str]:
        try:
            async for token in llm.stream(req.messages, req.temperature):
                if await request.is_disconnected():  # 客户端断开就停止生成，节省模型费用
                    log.info("客户端已断开，停止生成 request_id=%s", rid)
                    return
                yield f"data: {json.dumps({'delta': token}, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"
        except Exception:
            log.exception("生成失败 request_id=%s", rid)
            # 响应头已经发出，不能再改状态码，只能在流里发送错误事件
            yield f"event: error\ndata: {json.dumps({'message': '生成失败', 'request_id': rid}, ensure_ascii=False)}\n\n"

    return StreamingResponse(events(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.get("/boom")
async def boom():
    raise RuntimeError("数据库连接串 postgres://secret@...")  # 演示：内部细节不会泄露给调用方


if __name__ == "__main__":
    from fastapi.testclient import TestClient

    with TestClient(app) as c:
        r = c.post("/v1/chat", headers={"X-API-Key": "demo-key"},
                   json={"messages": [{"role": "user", "content": "年假怎么算"}], "stream": False})
        print(r.status_code, r.json())