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

from typing import Annotated, Literal

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

from .base import NonBlankString

Content = str | dict[str, JsonValue] | list[JsonValue]
Probability = Annotated[float, Field(ge=0, le=1)]


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)


class ScoreDecision(DecisionModel):
    type: Literal["score"] = "score"
    instructions: Content
    criteria: list[Content] = Field(min_length=2, max_length=10)


class NoulDecision(DecisionModel):
    type: Literal["noul"] = "noul"
    instructions: Content
    criteria: dict[Literal["true", "false"], Content] | None = None


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):
    type: Literal["noul"] = "noul"
    noul: Probability


DecisionResult = Annotated[
    ChoiceResult | ScoreResult | NoulResult, Field(discriminator="type")
]


class DecisionUsage(DecisionModel):
    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]
    model: str
    usage: DecisionUsage
    request_id: str | None = None