"""模型能力档案: 不同模型的参数差异都收敛到 yaml,加新模型不用改代码。"""
from __future__ import annotations

from dataclasses import dataclass, field, fields
from pathlib import Path
from typing import List

import yaml

from .llm_params import THINKING_TRANSPORTS

REASONING_REPLAY_POLICIES = {"none", "tool_calls", "all"}


@dataclass
class ModelCapabilities:
    model_id: str = ""
    family: str = ""
    variant: str = ""
    display_name: str = ""  # UI 展示用,如 "DeepSeek V4 Flash";空时前端 fallback 拼 family.variant

    # 上下文
    max_context: int = 128_000
    reliable_context: int = 64_000
    max_output: int = 4096

    # Tool calling
    parallel_tools: bool = False
    tool_calling_quality: str = "good"

    # 思考开关
    thinking_enabled: bool = False
    # none=不猜 provider 默认值；extra_body=显式发送 thinking.type，effort 同体透传。
    thinking_transport: str = "none"
    reasoning_effort_levels: List[str] = field(default_factory=list)
    default_reasoning_effort: str = ""
    # 历史 reasoning 发回模型的策略。持久化原始响应与 provider-bound 输入分离：
    # none=全部剥离；tool_calls=仅工具调用 assistant 保留；all=完整原样回传。
    reasoning_replay: str = "none"

    # 代码 / 沙盒
    code_quality: str = "good"
    enable_run_python: bool = False

    # 工程参数
    max_iterations: int = 120   # 单轮自主步数 backstop;空转防护见 loop 无进展熔断,不靠这个砍正经长任务
    optimal_temperature: float = 0.3

    # provider 特性
    prompt_caching: bool = False
    extended_thinking: bool = False

    # 计费兜底(CNY / million tokens)。provider / LiteLLM cost map 缺失时使用。
    input_cny_per_mtoken: float = 0.0
    output_cny_per_mtoken: float = 0.0
    # 前缀缓存命中价(DeepSeek 等自动缓存 prompt 前缀,命中部分按此价,通常 ~0.1x input)。
    # 0 = 不区分,缓存命中按 input 全价记(安全兜底,不会少记)。
    cache_hit_cny_per_mtoken: float = 0.0

    # API 接入
    api_base: str = ""
    api_key_env: str = ""

    @classmethod
    def load(cls, name: str, models_dir: Path) -> "ModelCapabilities":
        """name: '<family>.<variant>',如 'deepseek_v4.flash'。"""
        if "." in name:
            family, variant = name.split(".", 1)
        else:
            family, variant = name, "default"

        path = Path(models_dir) / f"{family}.yaml"
        if not path.exists():
            raise FileNotFoundError(f"模型档案不存在: {path}")

        data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
        variants = data.get("variants", {})
        if variant not in variants:
            raise ValueError(
                f"档案 {path} 没有 variant={variant};可选: {list(variants)}"
            )

        var = dict(variants[variant])
        valid_keys = {f.name for f in fields(cls)}
        kwargs = {k: v for k, v in var.items() if k in valid_keys}
        kwargs["family"] = data.get("family", family)
        kwargs["variant"] = variant
        caps = cls(**kwargs)
        if caps.thinking_transport not in THINKING_TRANSPORTS:
            raise ValueError(
                f"档案 {path} 的 thinking_transport={caps.thinking_transport!r} 无效;"
                f"可选: {sorted(THINKING_TRANSPORTS)}"
            )
        if caps.reasoning_replay not in REASONING_REPLAY_POLICIES:
            raise ValueError(
                f"档案 {path} 的 reasoning_replay={caps.reasoning_replay!r} 无效;"
                f"可选: {sorted(REASONING_REPLAY_POLICIES)}"
            )
        if caps.thinking_enabled and caps.thinking_transport == "none":
            raise ValueError(
                f"档案 {path} 开启 thinking 时必须声明可验证的 thinking_transport"
            )
        if (
            caps.default_reasoning_effort
            and caps.default_reasoning_effort not in caps.reasoning_effort_levels
        ):
            raise ValueError(
                f"档案 {path} 的 default_reasoning_effort="
                f"{caps.default_reasoning_effort!r} 不在 reasoning_effort_levels 中"
            )
        return caps