# -*- 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 json
import logging

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

from ... import config
from ...services import compute_catalog, compute_client
from ..dependencies import get_db
from ...infrastructure.repositories import Database

logger = logging.getLogger("relay")

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


def _token_fp(token: str) -> str:
    """令牌指纹（前 6 位 + 长度）：日志可定位但不可还原。"""
    return f"{token[:6]}…({len(token)})" if token else "(empty)"


def _client_ip(request: Request) -> str:
    fwd = request.headers.get("x-forwarded-for", "")
    return fwd.split(",")[0].strip() if fwd else (request.client.host if request.client else "")


def _is_jwt(token: str) -> bool:
    """平台 JWT 形如 header.payload.signature（两处 '.'）；引擎 PAT 不含 '.'。"""
    return token.count(".") == 2


def _strip_sk(token: str) -> str:
    """剥掉 OpenAI 惯例的 sk- 前缀：引擎 /v1 两种都认，但 /api/* 只认裸 key。"""
    return token[3:] if token.startswith("sk-") else token


async def _user_engine_pat(token: str) -> str:
    """把平台 JWT 解析为该用户的引擎消费令牌（PAT）。

    取该用户引擎名下第一枚有效令牌；没有则现场签发一枚，保证「登录用户=引擎用户」。
    """
    from ...jwt import decode_access_token

    claims = decode_access_token(token)
    username = (claims or {}).get("username") or ""
    if not username:
        raise HTTPException(status_code=401, detail="登录态无效，请重新登录")
    await compute_client.ensure_user(username)
    items = await compute_client.list_user_tokens(username)
    for it in items or []:
        key = it.get("key") or it.get("token")
        if key:
            return key
    issued = await compute_client.issue_user_token(username, name="平台中继令牌")
    key = issued.get("key") or issued.get("token") or ""
    if not key:
        raise HTTPException(status_code=502, detail="算力令牌签发失败，请联系运营方")
    logger.info("[relay] jwt→pat username=%s issued_new=%s", username, True)
    return key


async def _auth_headers(request: Request) -> dict[str, str]:
    """模型调用鉴权：把请求归到「真实用户令牌」，由引擎按该用户余额计量。

    - 客户端携带引擎 PAT：直接透传（compute 校验令牌与余额，无效即 401）。
    - 客户端携带平台 JWT（桌面端/网页端登录态）：解析用户 → 归户到其引擎 PAT。
    - 二者皆无：**拒绝匿名**——不回落 COMPUTE_RELAY_TOKEN/COMPUTE_ADMIN_TOKEN，
      否则「用户无余额仍能免费使用」。
    """
    client_auth = (request.headers.get("Authorization", "") or "").strip()
    auth_src = "Authorization"
    if not client_auth:
        # 兼容部分客户端/网关：key 也可能放在 x-api-key / api-key / 查询参数里
        for alt in ("x-api-key", "api-key", "x-goog-api-key"):
            v = (request.headers.get(alt, "") or "").strip()
            if v:
                client_auth = v
                auth_src = alt
                break
        else:
            for q in ("key", "api_key", "token"):
                v = (request.query_params.get(q, "") or "").strip()
                if v:
                    client_auth = v
                    auth_src = f"query:{q}"
                    break
    if not client_auth:
        logger.warning("[relay] 401 no-token %s %s ip=%s hdrs(auth=%s,apikey=%s,xapikey=%s)",
                       request.method, request.url.path, _client_ip(request),
                       bool(request.headers.get("authorization")),
                       bool(request.headers.get("api-key")),
                       bool(request.headers.get("x-api-key")))
        raise HTTPException(
            status_code=401,
            detail="缺少算力令牌：请携带 Bearer <用户算力令牌> 或登录态调用",
        )
    token = client_auth[7:].strip() if client_auth.lower().startswith("bearer ") else client_auth
    if not token:
        raise HTTPException(status_code=401, detail="算力令牌为空")
    kind = "jwt" if _is_jwt(token) else "pat"
    logger.info("[relay] auth %s %s ip=%s src=%s kind=%s token=%s",
                request.method, request.url.path, _client_ip(request), auth_src, kind, _token_fp(token))
    if _is_jwt(token):
        try:
            pat = await _user_engine_pat(token)
        except HTTPException:
            raise
        except Exception as exc:  # noqa: BLE001
            raise HTTPException(status_code=502, detail=f"算力引擎对接失败: {exc}") from exc
        return {"Authorization": f"Bearer {pat}"}
    return {"Authorization": f"Bearer {_strip_sk(token)}"}


@router.post("/chat/completions")
async def relay_chat_completions(request: Request, db: Database = Depends(get_db)):
    """把 chat/completions 转发到 compute-engine，流式 SSE 原样透传。

    平台 JWT 登录态调用：转发后按实际用量记账（先企业分配余额、后个人余额）。
    引擎 PAT 调用（第三方 OpenAI 客户端）：不记账，仅引擎计量。
    """
    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))
    model = str(body.get("model") or "")
    # 解析平台登录态 → user_id（仅 JWT 调用记账）
    billing_user_id = await _resolve_billing_user(request, db)
    headers = {**(await _auth_headers(request)), "Content-Type": "application/json"}
    logger.info("[relay] chat model=%s stream=%s bytes=%d bill=%s", model, stream, len(json.dumps(body)), bool(billing_user_id))

    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()
        # 按实际用量记账（best-effort，失败不影响响应）
        if billing_user_id and isinstance(payload, dict) and payload.get("usage"):
            await _bill_usage(db, billing_user_id, model, payload.get("usage"))
        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"]

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

    # 流式：透传 SSE 同时解析 usage chunk，结束后记账
    async def _stream_with_billing():
        usage: dict | None = None
        try:
            async for line in upstream.aiter_lines():
                yield (line + "\n").encode("utf-8")
                if line.startswith("data:"):
                    payload = line[5:].strip()
                    if payload and payload != "[DONE]":
                        try:
                            obj = json.loads(payload)
                            if isinstance(obj, dict) and obj.get("usage"):
                                usage = obj["usage"]
                        except Exception:  # noqa: BLE001
                            pass
        finally:
            try:
                await client.aclose()
            except Exception:  # noqa: BLE001
                pass
            if usage:
                await _bill_usage(db, billing_user_id, model, usage)

    return StreamingResponse(
        _stream_with_billing(),
        status_code=upstream.status_code,
        media_type="text/event-stream",
        headers=resp_headers,
    )


async def _resolve_billing_user(request: Request, db: Database) -> str:
    """解析平台 JWT 登录态 → 平台 user_id；非 JWT（引擎 PAT）返回空串不记账。"""
    client_auth = (request.headers.get("Authorization", "") or "").strip()
    if not client_auth:
        for alt in ("x-api-key", "api-key", "x-goog-api-key"):
            v = (request.headers.get(alt, "") or "").strip()
            if v:
                client_auth = v
                break
        else:
            for q in ("key", "api_key", "token"):
                v = (request.query_params.get(q, "") or "").strip()
                if v:
                    client_auth = v
                    break
    token = client_auth[7:].strip() if client_auth.lower().startswith("bearer ") else client_auth
    if not token or not _is_jwt(token):
        return ""
    try:
        from ...jwt import decode_access_token
        claims = decode_access_token(token)
        username = (claims or {}).get("username") or ""
        if not username:
            return ""
        user = await db.users.get_by_username(username)
        return (user or {}).get("id", "") or ""
    except Exception:  # noqa: BLE001
        return ""


async def _bill_usage(db: Database, user_id: str, model: str, usage: dict) -> None:
    """按实际用量记账：先企业分配余额、后个人余额（best-effort，失败仅记日志）。"""
    try:
        from ...services.compute_pricing_service import deduct_usage_post
        await deduct_usage_post(
            db, user_id, model,
            int(usage.get("prompt_tokens") or 0),
            int(usage.get("completion_tokens") or 0),
        )
    except Exception as exc:  # noqa: BLE001
        logger.warning("[relay] billing failed user=%s model=%s err=%s", user_id, model, exc)


@router.get("/models")
async def relay_models(_auth: dict = Depends(_auth_headers)):
    """可用模型清单（OpenAI GET /v1/models 口径，需令牌/登录态）。"""
    ids = await compute_catalog.model_ids()
    return {
        "object": "list",
        "data": [{"id": m, "object": "model", "owned_by": "pineagents"} for m in ids],
    }


@router.get("/models/{model_id}")
async def relay_model_detail(model_id: str, _auth: dict = Depends(_auth_headers)):
    """单个模型详情（OpenAI GET /v1/models/{model} 口径）。"""
    return {"id": model_id, "object": "model", "owned_by": "pineagents"}


# ---------------------------------------------------------------------------
# 余额（OpenAI billing 口径：多数客户端如 Cherry Studio/NextChat 用这两个端点展示余额）
# 口径：引擎 quota 单位 = 1e6 微元/元；subscription.hard_limit_usd = 总额度（元数），
# usage.total_usage = 已用（「美分」= 元 × 100），客户端显示 余额 = hard_limit - total_usage/100。
# ---------------------------------------------------------------------------
def _quota_per_unit() -> float:
    return 1_000_000.0


async def _engine_self(request: Request) -> dict:
    """按调用方令牌归户，查引擎 /api/user/self（余额/已用）。"""
    pat = (await _auth_headers(request))["Authorization"].split(" ", 1)[1]
    if _is_jwt(pat):
        raise HTTPException(status_code=401, detail="余额查询请携带用户算力令牌")
    pat = _strip_sk(pat)
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            r = await client.get(
                f"{config.COMPUTE_BASE_URL.rstrip('/')}/api/user/self",
                headers={"Authorization": f"Bearer {pat}"},
            )
    except Exception as exc:  # noqa: BLE001
        logger.warning("[relay] billing self 查询失败 token=%s err=%s", _token_fp(pat), exc)
        raise HTTPException(status_code=502, detail=f"算力引擎对接失败: {exc}") from exc
    data = (r.json() or {}).get("data") or {}
    # 引擎对无效令牌也返回 200 + data.error（如「无效令牌」），必须显式暴露而非静默 0
    if data.get("error"):
        logger.warning("[relay] billing self 无效令牌 token=%s engine=%s", _token_fp(pat), data.get("error"))
        raise HTTPException(status_code=401, detail=f"算力令牌无效：{data.get('error')}")
    if r.status_code != 200:
        logger.warning("[relay] billing self 非200 status=%s token=%s", r.status_code, _token_fp(pat))
        raise HTTPException(status_code=502, detail="算力引擎余额查询失败")
    logger.info("[relay] billing quota=%s used=%s", data.get("quota"), data.get("used_quota"))
    return data


@router.get("/dashboard/billing/subscription")
async def relay_billing_subscription(request: Request):
    self_data = await _engine_self(request)
    quota = float(self_data.get("quota") or 0)
    total = quota / _quota_per_unit()
    return {
        "object": "billing_subscription",
        "has_payment_method": True,
        "soft_limit_usd": total,
        "hard_limit_usd": total,
        "system_hard_limit_usd": total,
        "access_until": 0,
    }


@router.get("/dashboard/billing/usage")
async def relay_billing_usage(request: Request, date: str = ""):
    self_data = await _engine_self(request)
    used = float(self_data.get("used_quota") or 0)
    return {
        "object": "list",
        "total_usage": used / _quota_per_unit() * 100,  # 美分口径（元×100）
        "daily_costs": [],
    }


# ---------------------------------------------------------------------------
# 通用透传（OpenAI 规范其余端点：completions/embeddings/images/audio/moderations/
# responses/realtime… 以及 /v1 下任何路径），流式 SSE 原样透传，鉴权与 chat 同规。
# ---------------------------------------------------------------------------
@router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
async def relay_catch_all(request: Request, path: str):
    engine_url = f"{config.COMPUTE_BASE_URL.rstrip('/')}/v1/{path}"
    if str(request.url.query):
        engine_url = f"{engine_url}?{request.url.query}"

    headers = {**(await _auth_headers(request)), "Content-Type": request.headers.get("content-type", "application/json")}
    body = await request.body()

    # 流式判定：JSON body 里 stream=true（SSE 原样透传）
    wants_stream = False
    if body:
        try:
            wants_stream = bool(json.loads(body).get("stream"))
        except Exception:  # noqa: BLE001
            wants_stream = False

    client = httpx.AsyncClient(timeout=None)
    upstream_request = client.build_request(request.method, engine_url, content=body, headers=headers)
    logger.info("[relay] pass %s /%s → %s stream=%s bytes=%d", request.method, path, engine_url.split('?')[0], wants_stream, len(body))

    if wants_stream:
        try:
            upstream = await client.send(upstream_request, stream=True)
        except Exception:  # noqa: BLE001
            await client.aclose()
            logger.error("[relay] pass %s /%s 上游连接失败", request.method, path)
            raise HTTPException(status_code=502, detail="Upstream relay failed") from None
        logger.info("[relay] pass %s /%s → upstream %s (stream)", request.method, path, upstream.status_code)
        resp_headers = {k: upstream.headers[k] for k in ("content-type", "x-request-id") if k in upstream.headers}
        return StreamingResponse(
            upstream.aiter_raw(), status_code=upstream.status_code,
            media_type=upstream.headers.get("content-type", "application/json"),
            headers=resp_headers, background=BackgroundTask(client.aclose),
        )

    try:
        upstream = await client.send(upstream_request)
        payload = upstream.content
    except Exception:  # noqa: BLE001
        await client.aclose()
        logger.error("[relay] pass %s /%s 上游连接失败", request.method, path)
        raise HTTPException(status_code=502, detail="Upstream relay failed") from None
    await client.aclose()
    if upstream.status_code >= 400:
        logger.warning("[relay] pass %s /%s → upstream %s body=%s",
                       request.method, path, upstream.status_code, payload[:300])
    return Response(
        content=payload, status_code=upstream.status_code,
        media_type=upstream.headers.get("content-type", "application/json"),
    )