"""Token 消耗能力开关 — 全局 + per-user budget.yaml。"""

from __future__ import annotations

from copy import deepcopy
from pathlib import Path
from typing import Any, Literal

import yaml

from niannian.token.paths import token_budget_path, user_token_budget_path
from niannian.users.paths import user_root

ChatLlmLevel = Literal["off", "light", "heavy"]
CHAT_LLM_LEVELS: frozenset[str] = frozenset({"off", "light", "heavy"})

ChatMemoryMode = Literal["off", "manual_only", "batch"]
CHAT_MEMORY_MODES: frozenset[str] = frozenset({"off", "manual_only", "batch"})

ChatStickerLevel = Literal["off", "light", "heavy"]
CHAT_STICKER_LEVELS: frozenset[str] = frozenset({"off", "light", "heavy"})

ProactiveChatBurst = Literal["once", "few", "unlimited"]
PROACTIVE_CHAT_BURSTS: frozenset[str] = frozenset({"once", "few", "unlimited"})

ProactiveDisturbLock = Literal["auto", "quiet"]
PROACTIVE_DISTURB_LOCKS: frozenset[str] = frozenset({"auto", "quiet"})

DEFAULT_BOOL_BUDGET: dict[str, bool] = {
    "virtual_persona_enabled": True,
    "virtual_persona_enrich_enabled": True,
    "virtual_persona_avatar_fetch_enabled": True,
    "virtual_persona_sticker_search_enabled": True,
    "proactive_chat_enabled": False,
    "dream_consolidation_enabled": True,
    "web_learning_enabled": True,
    "cultural_learning_enabled": True,
    "embedding_recall_enabled": True,
}

DEFAULT_LEVEL_BUDGET: dict[str, str] = {
    "chat_llm_level": "heavy",
    "chat_memory_mode": "manual_only",
    "chat_sticker_level": "light",
    "proactive_chat_burst": "few",
    "proactive_disturb_lock": "auto",
}

LEVEL_KEY_VALIDATORS: dict[str, frozenset[str]] = {
    "chat_llm_level": CHAT_LLM_LEVELS,
    "chat_memory_mode": CHAT_MEMORY_MODES,
    "chat_sticker_level": CHAT_STICKER_LEVELS,
    "proactive_chat_burst": PROACTIVE_CHAT_BURSTS,
    "proactive_disturb_lock": PROACTIVE_DISTURB_LOCKS,
}

BUDGET_BOOL_KEYS = frozenset(DEFAULT_BOOL_BUDGET.keys())
BUDGET_LEVEL_KEYS = frozenset(DEFAULT_LEVEL_BUDGET.keys())
BUDGET_KEYS = BUDGET_BOOL_KEYS | BUDGET_LEVEL_KEYS


def _merge_yaml_into(out: dict[str, Any], raw: dict[str, Any]) -> None:
    for key in BUDGET_BOOL_KEYS:
        if key in raw:
            out[key] = bool(raw[key])
    for key in BUDGET_LEVEL_KEYS:
        if key in raw:
            level = str(raw[key] or "").strip().lower()
            if level in LEVEL_KEY_VALIDATORS.get(key, frozenset()):
                out[key] = level


def load_token_budget(home: Path, user_id: str = "pinn") -> dict[str, Any]:
    """合并默认 · 全局 budget · 用户 budget（用户优先）。"""
    out: dict[str, Any] = {**DEFAULT_BOOL_BUDGET, **DEFAULT_LEVEL_BUDGET}
    global_path = token_budget_path(home)
    if global_path.is_file():
        with global_path.open(encoding="utf-8") as f:
            raw = yaml.safe_load(f) or {}
        if isinstance(raw, dict):
            _merge_yaml_into(out, raw)
    uid = str(user_id or "pinn").strip() or "pinn"
    user_path = user_token_budget_path(home, uid)
    if user_path.is_file():
        with user_path.open(encoding="utf-8") as f:
            raw = yaml.safe_load(f) or {}
        if isinstance(raw, dict):
            _merge_yaml_into(out, raw)
    return out


def save_token_budget(home: Path, patch: dict[str, Any], *, user_id: str = "pinn") -> dict[str, Any]:
    """写入 users/{user_id}/token/budget.yaml（普通用户各自偏好）。"""
    current = load_token_budget(home, user_id)
    for key, val in patch.items():
        if key in BUDGET_BOOL_KEYS:
            current[key] = bool(val)
        elif key in BUDGET_LEVEL_KEYS:
            level = str(val or "").strip().lower()
            if level in LEVEL_KEY_VALIDATORS.get(key, frozenset()):
                current[key] = level
    uid = str(user_id or "pinn").strip() or "pinn"
    path = user_token_budget_path(home, uid)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(yaml.safe_dump(current, allow_unicode=True, sort_keys=True), encoding="utf-8")
    return current


def budget_allows(budget: dict[str, Any], key: str) -> bool:
    if key in BUDGET_BOOL_KEYS:
        return bool(budget.get(key, DEFAULT_BOOL_BUDGET.get(key, True)))
    return True


def resolve_chat_llm_level(budget: dict[str, Any]) -> ChatLlmLevel:
    raw = str(budget.get("chat_llm_level") or DEFAULT_LEVEL_BUDGET["chat_llm_level"]).lower()
    if raw in CHAT_LLM_LEVELS:
        return raw  # type: ignore[return-value]
    return "heavy"


def resolve_chat_memory_mode(budget: dict[str, Any]) -> ChatMemoryMode:
    raw = str(budget.get("chat_memory_mode") or DEFAULT_LEVEL_BUDGET["chat_memory_mode"]).lower()
    if raw in CHAT_MEMORY_MODES:
        return raw  # type: ignore[return-value]
    return "manual_only"


def chat_llm_allows(budget: dict[str, Any], tier: Literal["reply", "auxiliary"]) -> bool:
    """reply=对话主回复 · auxiliary=ingest判定/记忆查询/模式LLM等附加调用。"""
    level = resolve_chat_llm_level(budget)
    if tier == "reply":
        return level in ("light", "heavy")
    return level == "heavy"


def resolve_chat_sticker_level(budget: dict[str, Any]) -> ChatStickerLevel:
    raw = str(budget.get("chat_sticker_level") or DEFAULT_LEVEL_BUDGET["chat_sticker_level"]).lower()
    if raw in CHAT_STICKER_LEVELS:
        return raw  # type: ignore[return-value]
    return "light"


def chat_sticker_allows(
    budget: dict[str, Any],
    tier: Literal["user", "assistant", "learn"],
) -> bool:
    """user=用户发表情 · assistant=念念主动表情 · learn=表情包 label 入库。"""
    level = resolve_chat_sticker_level(budget)
    if level == "off":
        return False
    if tier == "learn":
        return level == "heavy"
    return True


def resolve_proactive_chat_burst(budget: dict[str, Any]) -> ProactiveChatBurst:
    raw = str(
        budget.get("proactive_chat_burst") or DEFAULT_LEVEL_BUDGET["proactive_chat_burst"]
    ).lower()
    if raw in PROACTIVE_CHAT_BURSTS:
        return raw  # type: ignore[return-value]
    return "few"


def resolve_proactive_disturb_lock(budget: dict[str, Any]) -> ProactiveDisturbLock:
    raw = str(
        budget.get("proactive_disturb_lock") or DEFAULT_LEVEL_BUDGET["proactive_disturb_lock"]
    ).lower()
    if raw in PROACTIVE_DISTURB_LOCKS:
        return raw  # type: ignore[return-value]
    return "auto"
