"""代理面路由（T-P4）：学生面对话 + 管理面 CRUD。

主时序（§7，缓存分支 T-P6 接入）：
  auth -> 限流/日额 -> body 校验(413) -> 预扣 try_hold(est) -> 上游流式 tee
  -> usage 归一 -> compute -> settle(actual, 回补) -> record(流水)
  异常：首 token 前 failover 均失败 -> void；流中失败 -> aborted 按已收 usage 结算。

管理面：X-Admin-Key（hmac.compare_digest；未配置仅 loopback）。
"""
from __future__ import annotations

import asyncio
import json
import time
import uuid
from typing import Any, Dict, Optional

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, StreamingResponse

from gateway.proxy.auth import authenticate, verify_admin
from gateway.proxy.config import ProxyConfig
from gateway.proxy.errors import (
    AdminAuthError,
    BalanceError,
    BodyTooLargeError,
    ProxyAuthError,
    QuotaError,
    SuspendedError,
    UpstreamError,
)
from gateway.proxy.pricing import compute
from gateway.proxy.upstream import UpstreamAborted, filter_usage_chunk, stream as upstream_stream


def install_error_handlers(app) -> None:
    """把 ProxyError 家族映射为 §5.3 错误码（401/402/403/413/429/502）。"""
    from gateway.proxy.errors import ProxyError

    @app.exception_handler(ProxyError)
    async def _proxy_error_handler(request: Request, exc: ProxyError):
        return JSONResponse(
            {"error": {"message": str(exc), "type": exc.code}},
            status_code=exc.status_code,
            headers={"WWW-Authenticate": "Bearer"} if exc.status_code == 401 else None)


def build_proxy_router(cfg: ProxyConfig, pool) -> APIRouter:
    """组装代理面路由（唯一组装点）。"""
    router = APIRouter(prefix="/proxy")
    ledger = __import__("gateway.proxy.ledger", fromlist=["Ledger"]).Ledger.init_db(cfg.db_path)
    # ---------------- 学生面 ----------------
    @router.get("/v1/models", tags=["proxy"])
    async def list_models():
        """学生面：池内允许代理的模型列表（OpenAI /models 形状）。"""
        entries = pool.list().get("entries", [])
        seen, data = set(), []
        for e in entries:
            if not e.get("enabled") or e.get("backend") in ("mock",):
                continue
            mid = e.get("model") or ""
            if mid and mid not in seen:
                seen.add(mid)
                data.append({"id": mid, "object": "model", "owned_by": "campus-proxy"})
        return {"object": "list", "data": data}

    @router.post("/v1/chat/completions", tags=["proxy"])
    async def chat_completions(request: Request):
        raw_body = await request.body()
        if len(raw_body) > cfg.max_body_chars:
            raise BodyTooLargeError(f"请求体超过 {cfg.max_body_chars} 字符")
        try:
            body = json.loads(raw_body or b"{}")
        except json.JSONDecodeError:
            return JSONResponse({"error": {"message": "请求体不是合法 JSON"}}, status_code=400)

        ctx = authenticate(request.headers.get("authorization", ""), ledger,
                           _limits(), time.time())
        if not _limits().acquire_slot(ctx["key_id"]):
            raise QuotaError("并发请求已达该 key 上限")
        try:
            return await _run_chat(body, dict(request.headers), ctx, cfg, ledger, pool)
        finally:
            _limits().release_slot(ctx["key_id"])

    # ---------------- 管理面 ----------------
    def _guard_admin(request: Request) -> None:
        host = request.client.host if request.client else ""
        if not verify_admin(request.headers.get("x-admin-key", ""), cfg.admin_key, host):
            raise AdminAuthError("管理面鉴权失败")

    @router.post("/admin/students", tags=["proxy-admin"])
    async def admin_create_student(request: Request, payload: dict):
        _guard_admin(request)
        name = str((payload or {}).get("name") or "").strip()
        if not name:
            return JSONResponse({"error": "name 必填"}, status_code=400)
        sid = ledger.upsert_student(
            name, klass=str((payload or {}).get("class") or ""),
            balance_yuan=float((payload or {}).get("balance_yuan", 0) or 0),
            daily_cap_yuan=float((payload or {}).get("daily_cap_yuan", 5) or 5))
        return {"student_id": sid, **(ledger.get_student(sid) or {})}

    @router.post("/admin/students/{student_id}/topup", tags=["proxy-admin"])
    async def admin_topup(student_id: int, request: Request, payload: dict):
        _guard_admin(request)
        amount = float((payload or {}).get("amount_yuan", 0) or 0)
        new_balance = ledger.topup(student_id, amount)
        if new_balance is None:
            return JSONResponse({"error": f"学生不存在: {student_id}"}, status_code=404)
        return {"ok": True, "balance_milli": new_balance}

    @router.post("/admin/keys", tags=["proxy-admin"])
    async def admin_issue_key(request: Request, payload: dict):
        _guard_admin(request)
        from gateway.proxy.auth import issue_key
        student_id = int((payload or {}).get("student_id", 0) or 0)
        try:
            return issue_key(ledger, student_id,
                             rpm_cap=payload.get("rpm_cap"),
                             day_cap_req=payload.get("day_cap_req"))
        except ProxyAuthError as e:
            return JSONResponse({"error": str(e)}, status_code=404)

    @router.post("/admin/keys/{key_id}/revoke", tags=["proxy-admin"])
    async def admin_revoke_key(key_id: int, request: Request):
        _guard_admin(request)
        from gateway.proxy.auth import _AUTH_SINGLETON
        ok = ledger.revoke_key(key_id)
        return {"ok": ok}

    return router


def _limits():
    from gateway.proxy.auth import _AUTH_SINGLETON_LIMITS
    return _AUTH_SINGLETON_LIMITS


def _estimate_hold_milli(body: dict, cfg: ProxyConfig, model: str) -> int:
    """预扣估算（宁可高估，D-P11）：in 按字符/3，out 按 min(max_tokens,4096)。

    只能用已配置模型的峰值单价；未配置模型按 default 桶兜底价 or 最高价。
    """
    price = cfg.price(model)
    if price is None:
        prices = cfg.model_prices.values()
        if not prices:
            return 1
        in_p = max(p.in_miss for p in prices)
        hit_p = max(p.in_hit for p in prices)
        out_p = max(p.out for p in prices)
    else:
        in_p, hit_p, out_p = price.in_miss, price.in_hit, price.out
    prompt_chars = sum(len(str(m.get("content") or ""))
                       for m in (body.get("messages") or []))
    est_in = prompt_chars // 3
    est_out = int(body.get("max_tokens") or 1024)
    est_out = min(est_out, 4096)
    cost = (est_in * in_p + est_out * out_p) / 1_000_000
    return max(1, int(round(cost)))          # 保守：全部按未命中价


def _resolve_entry(pool, model: str, cfg: ProxyConfig) -> Optional[Dict[str, Any]]:
    """按模型名取池条目（仅启用的真实后端）。"""
    e = pool.find_by_model(model)
    if e and e.get("enabled") and e.get("backend") not in ("mock",) and e.get("base_url"):
        return e
    return None


async def _run_chat(body: dict, headers: Dict[str, str], ctx: Dict[str, Any],
                    cfg: ProxyConfig, ledger, pool):
    model = str(body.get("model") or "")
    entry = _resolve_entry(pool, model, cfg)
    if entry is None:
        return JSONResponse({"error": {"message": f"模型不可用: {model}",
                                       "type": "invalid_request_error"}},
                            status_code=400)
    request_id = "px" + uuid.uuid4().hex[:12]
    ts = time.time()
    client_wants_usage = bool(body.get("stream_options", {}).get("include_usage")) \
        if isinstance(body.get("stream_options"), dict) else False
    is_stream = bool(body.get("stream"))

    est = _estimate_hold_milli(body, cfg, model)
    if not await asyncio.to_thread(
            ledger.try_hold, request_id, ctx["key_id"], ctx["student_id"],
            model, "default", est, ts):
        raise BalanceError("余额或当日额度不足")

    t0 = time.perf_counter()
    sink: Dict[str, Any] = {}
    try:
        if is_stream:
            return await _stream_response(body, entry, sink, headers, client_wants_usage,
                                          request_id, ctx, cfg, ledger, model, est, t0)
        return await _json_response(body, entry, sink, request_id, ctx, cfg,
                                    ledger, model, est, t0)
    except UpstreamAborted as e:
        # 流中失败：按已收 usage 结算（无 usage 按字符估算），不缓存（D-P4）
        usage = sink.get("usage") or _estimate_usage_from_sink(sink)
        br = compute(usage, model, ts, cfg)
        await asyncio.to_thread(
            ledger.settle, request_id, br["charged_milli"],
            in_miss_tok=usage.get("in_miss", 0), in_hit_tok=usage.get("in_hit", 0),
            out_tok=usage.get("out", 0), upstream_cost_milli=br["upstream_cost_milli"],
            total_ms=int((time.perf_counter() - t0) * 1000), status="aborted")
        return JSONResponse({"error": {"message": f"上游流中断: {e}",
                                       "type": "aborted"}}, status_code=502)
    except UpstreamError as e:
        # 首 token 前失败（failover 均失败）：全额退款
        await asyncio.to_thread(ledger.void, request_id, status="error")
        return JSONResponse({"error": {"message": str(e), "type": "upstream_error"}},
                            status_code=502)
    except Exception as e:  # noqa: BLE001
        await asyncio.to_thread(ledger.void, request_id, status="error")
        raise


def _estimate_usage_from_sink(sink: Dict[str, Any]) -> Dict[str, int]:
    text = sink.get("text", "")
    return {"in_miss": 0, "in_hit": 0, "out": len(text) // 4}


async def _stream_response(body, entry, sink, headers, client_wants_usage,
                           request_id, ctx, cfg, ledger, model, est, t0):
    usage = {"in_miss": 0, "in_hit": 0, "out": 0}

    async def gen():
        collected = []
        chunk_id = f"chatcmpl-{request_id}"
        created = int(time.time())
        try:
            async for raw_bytes in upstream_stream(body, entry, sink, [entry]):
                line = raw_bytes.decode("utf-8").strip()
                if not line:
                    continue
                for sub in line.split("\n\n"):
                    if not sub:
                        continue
                    filtered = filter_usage_chunk(sub, client_wants_usage)
                    if filtered is None:
                        continue
                    # 同构补齐（§5.1）：确保 OpenAI chunk 形状（object/created/id/model）
                    out_line = filtered
                    if filtered.startswith("data:") and "[DONE]" not in filtered:
                        try:
                            obj = json.loads(filtered[5:].strip())
                            obj.setdefault("object", "chat.completion.chunk")
                            obj.setdefault("id", chunk_id)
                            obj.setdefault("created", created)
                            obj.setdefault("model", model)
                            delta = (obj.get("choices") or [{}])[0].get("delta") or {}
                            collected.append(str(delta.get("content") or ""))
                            out_line = "data: " + json.dumps(obj, ensure_ascii=False)
                        except (json.JSONDecodeError, IndexError):
                            pass
                    yield (out_line + "\n\n").encode("utf-8")
            sink["text"] = "".join(collected)
        finally:
            u = sink.get("usage") or {"in_miss": 0, "in_hit": 0, "out": 0}
            br = compute(u, model, time.time(), cfg)
            await asyncio.to_thread(
                ledger.settle, request_id, br["charged_milli"],
                in_miss_tok=u.get("in_miss", 0), in_hit_tok=u.get("in_hit", 0),
                out_tok=u.get("out", 0), upstream_cost_milli=br["upstream_cost_milli"],
                ttfb_ms=sink.get("ttfb_ms"),
                total_ms=int((time.perf_counter() - t0) * 1000),
                status="ok")

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache",
                                      "X-Request-Id": request_id})


async def _json_response(body, entry, sink, request_id, ctx, cfg, ledger,
                         model, est, t0):
    parts = []
    async for raw_bytes in upstream_stream(body, entry, sink, [entry]):
        line = raw_bytes.decode("utf-8").strip()
        if not line:
            continue
        for sub in line.split("\n\n"):
            if sub.startswith("data:"):
                payload = sub[5:].strip()
                if payload == "[DONE]":
                    continue
                try:
                    obj = json.loads(payload)
                except json.JSONDecodeError:
                    continue
                delta = (obj.get("choices") or [{}])[0].get("delta") or {}
                if delta.get("content"):
                    parts.append(str(delta["content"]))
                if obj.get("finish_reason") or (obj.get("choices") or [{}])[0].get("finish_reason"):
                    sink.setdefault("finish_reason",
                                    (obj.get("choices") or [{}])[0].get("finish_reason"))
    usage = sink.get("usage") or {"in_miss": 0, "in_hit": 0,
                                  "out": len("".join(parts)) // 4}
    br = compute(usage, model, time.time(), cfg)
    total_ms = int((time.perf_counter() - t0) * 1000)
    await asyncio.to_thread(
        ledger.settle, request_id, br["charged_milli"],
        in_miss_tok=usage.get("in_miss", 0), in_hit_tok=usage.get("in_hit", 0),
        out_tok=usage.get("out", 0), upstream_cost_milli=br["upstream_cost_milli"],
        ttfb_ms=sink.get("ttfb_ms"), total_ms=total_ms, status="ok")
    return JSONResponse({
        "id": f"chatcmpl-{request_id}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [{"index": 0,
                     "message": {"role": "assistant",
                                 "content": "".join(parts)},
                     "finish_reason": sink.get("finish_reason") or "stop"}],
        "usage": {"prompt_tokens": usage.get("in_miss", 0) + usage.get("in_hit", 0),
                  "completion_tokens": usage.get("out", 0),
                  "total_tokens": usage.get("in_miss", 0) + usage.get("in_hit", 0)
                  + usage.get("out", 0)},
    }, headers={"X-Request-Id": request_id})