"""通用决策协议"""

from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator

from .base import NonBlankString

# 文本或结构化 JSON，用于传递上下文和判断说明
Content = str | dict[str, JsonValue] | list[JsonValue]
# 概率和置信度都限制在 0～1
Probability = Annotated[float, Field(ge=0, le=1)]
# 意图门控后的去向：明确的问答走问答节点，其余交给 screen agent
Route = Literal["qa", "screen"]


class DecisionModel(BaseModel):
    """公共校验规则：拒绝未知字段、隐式类型转换和非有限数值"""

    model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False)


class ChoiceDecision(DecisionModel):
    """选择题：从候选项中选择一个"""

    type: Literal["choice"] = "choice"
    instructions: Content  # 告诉模型要判断什么
    criteria: dict[NonBlankString, Content | None] = Field(min_length=1, max_length=255)
    # criteria 的键是选项，值是说明，不需要说明时填 None


class ScoreDecision(DecisionModel):
    """评分题：按从低到高排列的等级进行评分"""

    type: Literal["score"] = "score"
    instructions: Content
    criteria: list[Content] = Field(min_length=2, max_length=10)  # 等级从 0 开始编号


class NoulDecision(DecisionModel):
    """判断题：返回命题成立的概率，不直接返回布尔值"""

    type: Literal["noul"] = "noul"
    instructions: Content
    criteria: dict[Literal["true", "false"], Content] | None = (
        None  # 可选的成立与不成立说明
    )


# 根据 type 选择对应模型，避免混用选择题和评分题的字段
DecisionSchema = Annotated[
    ChoiceDecision | ScoreDecision | NoulDecision, Field(discriminator="type")
]


class DecisionRequest(DecisionModel):
    """一次调用的输入，可同时提交多个独立判断"""

    state: Content  # 所有判断共享的业务上下文
    decisions: dict[NonBlankString, DecisionSchema] = Field(
        min_length=1
    )  # 键用于对应返回结果
    model: NonBlankString | None = None  # 不传时使用客户端默认模型
    # 官方扩展字段入口，不允许覆盖已经校验的协议字段
    extra_body: dict[str, JsonValue] = Field(default_factory=dict)

    @model_validator(mode="after")
    def check_extra_body(self):
        if {"state", "questions", "model"} & self.extra_body.keys():
            raise ValueError("extra_body 不能覆盖 state、questions 或 model")
        return self


class ChoiceResult(DecisionModel):
    """选择结果及所有候选项的概率分布"""

    type: Literal["choice"] = "choice"
    choice: str  # 选中的候选项
    confidence: Probability  # 分布的集中程度，不等于选中项的概率
    probabilities: dict[str, Probability]  # 选项 → 概率


class ScoreResult(DecisionModel):
    """评分结果，保留等级说明和概率分布"""

    type: Literal["score"] = "score"
    score: float = Field(ge=0)  # 按等级概率加权的分数，可能是小数
    confidence: Probability
    legend: dict[str, Content]  # 等级编号字符串 → 等级说明
    probabilities: dict[str, Probability]  # 等级编号字符串 → 概率


class NoulResult(DecisionModel):
    """成立概率：接近 1 表示是，接近 0 表示否，接近 0.5 表示不确定"""

    type: Literal["noul"] = "noul"
    noul: Probability  # 没有单独的 confidence，业务自行决定判断阈值


# 返回结果也按 type 区分，读取前先确认结果类型
DecisionResult = Annotated[
    ChoiceResult | ScoreResult | NoulResult, Field(discriminator="type")
]


class DecisionUsage(DecisionModel):
    """本次调用的 token 用量，None 表示服务端未提供"""

    input_tokens: int | None = Field(default=None, ge=0)
    output_tokens: int | None = Field(default=None, ge=0)


class DecisionResponse(DecisionModel):
    """统一返回结构，包含判断结果和调用信息"""

    results: dict[str, DecisionResult]  # 键与请求中的 decisions 一一对应
    model: str  # 实际执行判断的模型
    usage: DecisionUsage
    request_id: str | None = None  # 服务端请求标识，方便排查问题