# -*- coding: utf-8 -*-
"""OpenAI 兼容模型中转端点。

PineAgents(8088) 的 ``pineagents`` provider 把 base_url 指向本服务的 ``/v1``，
请求进来后丢弃客户端任何 ``Authorization``，统一替换为服务端凭据。

模型调用统一经 **compute-engine**（loopback :3000）中转计量（引擎持有渠道 key，
前端不直连）；走 ``PINEAGENTS_COMPUTE_ADMIN_TOKEN`` 管理令牌。流式 SSE 原样透传。
"""
from __future__ import annotations

from typing import Any

import httpx
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask

from ... import config

router = APIRouter(prefix="/v1", tags=["relay"])

# 模型清单（供 /v1/models 使用，与 8088 侧保持一致）
_ENGINE_MODEL_IDS = [
    "deepseek-chat",
    "deepseek-reasoner",
    "deepseek-v4-flash",
    "deepseek-v4-pro",
]


def _auth_headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {config.COMPUTE_ADMIN_TOKEN}"}


@router.post("/chat/completions")
async def relay_chat_completions(request: Request):
    """把 chat/completions 转发到 compute-engine，流式 SSE 原样透传。"""
    try:
        body: dict[str, Any] = await request.json()
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=400, detail="Invalid JSON body") from exc

    chat_url = f"{config.COMPUTE_BASE_URL.rstrip('/')}/v1/chat/completions"
    stream = bool(body.get("stream", False))
    headers = {**_auth_headers(), "Content-Type": "application/json"}

    client = httpx.AsyncClient(timeout=None)
    upstream_request = client.build_request("POST", chat_url, json=body, headers=headers)

    if not stream:
        try:
            upstream = await client.send(upstream_request)
            payload = upstream.json() if upstream.content else None
        except Exception:  # noqa: BLE001
            await client.aclose()
            raise HTTPException(status_code=502, detail="Upstream relay failed") from None
        await client.aclose()
        return JSONResponse(status_code=upstream.status_code, content=payload)

    try:
        upstream = await client.send(upstream_request, stream=True)
    except Exception:  # noqa: BLE001
        await client.aclose()
        raise HTTPException(status_code=502, detail="Upstream relay failed") from None

    resp_headers: dict[str, str] = {}
    if "content-type" in upstream.headers:
        resp_headers["content-type"] = upstream.headers["content-type"]
    if "x-request-id" in upstream.headers:
        resp_headers["x-request-id"] = upstream.headers["x-request-id"]

    return StreamingResponse(
        upstream.aiter_raw(),
        status_code=upstream.status_code,
        media_type="text/event-stream",
        headers=resp_headers,
        background=BackgroundTask(client.aclose),
    )


@router.get("/models")
async def relay_models():
    """返回模型清单（供 OpenAIProvider.fetch_models 使用）。"""
    return {
        "object": "list",
        "data": [
            {"id": m, "object": "model", "owned_by": "pineagents"}
            for m in _ENGINE_MODEL_IDS
        ],
    }