"""Jev wire protocol (``POST /v1/systemone``) as pydantic models, 100% compatible with Jue."""

from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final, Literal, TypeAlias

from pydantic import BaseModel, ConfigDict, Field, JsonValue

JSONValue: TypeAlias = str | dict[str, JsonValue] | list[JsonValue]

MIN_CHOICE_OPTIONS: Final = 2
MAX_CHOICE_OPTIONS: Final = 255
MIN_SCORE_LEVELS: Final = 2
MAX_SCORE_LEVELS: Final = 10
MIN_QUESTIONS: Final = 1
MAX_QUESTIONS: Final = 256

NOUL_OPTION_KEYS: Final[tuple[str, str]] = ("false", "true")


class _Strict(BaseModel):
    model_config = ConfigDict(extra="forbid")


class NoulCriteria(_Strict):
    model_config = ConfigDict(serialize_by_alias=True)

    true_: JSONValue | None = Field(default=None, alias="true")
    false_: JSONValue | None = Field(default=None, alias="false")


class NoulQuestion(_Strict):
    type: Literal["noul"]
    instructions: JSONValue | None = None
    criteria: NoulCriteria | None = None


class ChoiceQuestion(_Strict):
    type: Literal["choice"]
    instructions: JSONValue | None = None
    criteria: dict[str, JSONValue | None]


class ScoreQuestion(_Strict):
    type: Literal["score"]
    instructions: JSONValue | None = None
    criteria: list[JSONValue]


Question = Annotated[NoulQuestion | ChoiceQuestion | ScoreQuestion, Field(discriminator="type")]


class JueOptions(_Strict):
    rounding: int | None = Field(default=None, ge=0, strict=True)
    calibration: str | None = None


class SystemOneRequest(_Strict):
    state: JSONValue
    model: str
    questions: dict[str, Question]
    jue_options: JueOptions | None = None


Probability = Annotated[float, Field(ge=0.0, le=1.0, allow_inf_nan=False)]


class NoulAnswer(_Strict):
    type: Literal["noul"]
    noul: Probability


class ChoiceAnswer(_Strict):
    type: Literal["choice"]
    choice: str
    probabilities: dict[str, Probability]
    confidence: Probability


class ScoreAnswer(_Strict):
    type: Literal["score"]
    score: float = Field(ge=0.0, allow_inf_nan=False)
    legend: dict[str, JSONValue]
    probabilities: dict[str, Probability]
    confidence: Probability


Answer = Annotated[NoulAnswer | ChoiceAnswer | ScoreAnswer, Field(discriminator="type")]


class Usage(_Strict):
    input_tokens: int = Field(ge=0)
    output_tokens: int = Field(ge=0)


class SystemOneResponse(_Strict):
    model: str = Field(pattern=r"^[a-zA-Z0-9._/-]+$", min_length=1, max_length=128)
    answers: dict[str, Answer]
    usage: Usage


ErrorType: TypeAlias = Literal[
    "invalid_api_key",
    "invalid_request",
    "too_many_options",
    "too_many_levels",
    "too_many_questions",
    "max_tokens_exceeded",
    "unknown_model",
]

ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType(
    {
        "invalid_api_key": 401,
        "invalid_request": 422,
        "too_many_options": 422,
        "too_many_levels": 422,
        "too_many_questions": 422,
        "max_tokens_exceeded": 422,
        "unknown_model": 422,
    }
)


class JueError(Exception):
    """A request-level failure with a stable ``error_type``; ``status`` is derived from it."""

    status: int
    error_type: ErrorType
    message: str

    def __init__(self, error_type: ErrorType, message: str) -> None:
        if error_type not in ERROR_STATUS:
            raise ValueError(f"unknown error_type {error_type!r}; expected one of {sorted(ERROR_STATUS)}")
        super().__init__(f"{error_type}: {message}")
        self.error_type = error_type
        self.status = ERROR_STATUS[error_type]
        self.message = message


def option_keys(question: Question) -> list[str]:
    """Answer keys in readout order: Choice criteria order, Noul ['false', 'true'], Score '0'..'K-1'."""
    if isinstance(question, ChoiceQuestion):
        return list(question.criteria)
    if isinstance(question, NoulQuestion):
        return list(NOUL_OPTION_KEYS)
    if isinstance(question, ScoreQuestion):
        return [str(level) for level in range(len(question.criteria))]
    raise TypeError(f"not a question model: {type(question).__name__}")


def check_question_limits(key: str, question: Question) -> None:
    over: ErrorType
    if isinstance(question, ChoiceQuestion):
        lo, hi, unit, over = MIN_CHOICE_OPTIONS, MAX_CHOICE_OPTIONS, "options", "too_many_options"
    elif isinstance(question, ScoreQuestion):
        lo, hi, unit, over = MIN_SCORE_LEVELS, MAX_SCORE_LEVELS, "levels", "too_many_levels"
    else:
        return
    n = len(question.criteria)
    if not lo <= n <= hi:
        message = f"{question.type} question {key!r} has {n} {unit}; allowed {lo}..{hi}"
        raise JueError(over if n > hi else "invalid_request", message)


def validate_limits(req: SystemOneRequest) -> None:
    n_questions = len(req.questions)
    if n_questions > MAX_QUESTIONS:
        raise JueError(
            "too_many_questions", f"request has {n_questions} questions; at most {MAX_QUESTIONS} are allowed"
        )
    if n_questions < MIN_QUESTIONS:
        raise JueError("invalid_request", f"request has no questions; at least {MIN_QUESTIONS} is required")
    for key, question in req.questions.items():
        check_question_limits(key, question)