"""Token 用量统计与消耗能力开关 — /v1/token/*。"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from fastapi import APIRouter, Query, Request
from pydantic import BaseModel, Field

from niannian.token.budget import BUDGET_KEYS, load_token_budget, save_token_budget
from niannian.tenant.context import resolve_request_context

if TYPE_CHECKING:
    from niannian.client import NianNian


class TokenBudgetPatchBody(BaseModel):
    virtual_persona_enabled: bool | None = None
    virtual_persona_enrich_enabled: bool | None = None
    virtual_persona_avatar_fetch_enabled: bool | None = None
    proactive_chat_enabled: bool | None = None
    dream_consolidation_enabled: bool | None = None
    web_learning_enabled: bool | None = None
    cultural_learning_enabled: bool | None = None
    embedding_recall_enabled: bool | None = None
    chat_llm_level: str | None = Field(default=None, pattern="^(off|light|heavy)$")
    chat_memory_mode: str | None = Field(default=None, pattern="^(off|manual_only|batch)$")
    chat_sticker_level: str | None = Field(default=None, pattern="^(off|light|heavy)$")
    proactive_chat_burst: str | None = Field(default=None, pattern="^(once|few|unlimited)$")
    proactive_disturb_lock: str | None = Field(default=None, pattern="^(auto|quiet)$")


def create_token_router(nn: NianNian) -> APIRouter:
    router = APIRouter(prefix="/v1/token", tags=["token"])

    @router.get("/usage")
    def get_token_usage(
        request: Request,
        days: int = Query(default=7, ge=1, le=7),
    ) -> dict[str, Any]:
        ctx = resolve_request_context(request, nn._config, home=nn._home)
        summary = nn.token_usage.summary(user_id=ctx.user_id, days=days)
        summary["user_id"] = ctx.user_id
        return summary

    @router.get("/budget")
    def get_token_budget(request: Request) -> dict[str, Any]:
        ctx = resolve_request_context(request, nn._config, home=nn._home)
        budget = load_token_budget(nn._home, ctx.user_id)
        return {"user_id": ctx.user_id, "budget": budget, "keys": sorted(BUDGET_KEYS)}

    @router.patch("/budget")
    def patch_token_budget(request: Request, body: TokenBudgetPatchBody) -> dict[str, Any]:
        ctx = resolve_request_context(request, nn._config, home=nn._home)
        patch = body.model_dump(exclude_unset=True)
        budget = save_token_budget(nn._home, patch, user_id=ctx.user_id)
        return {"user_id": ctx.user_id, "budget": budget}

    return router
