"""TypeSafe-compatible wire contract and deterministic probability projection."""

import json
import math
from typing import Annotated, Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

Content = str | dict[str, Any] | list[Any]


class QuestionBase(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    instructions: Content | None = None


class Choice(QuestionBase):
    type: Literal["choice"]
    criteria: dict[str, Content | None] = Field(min_length=1, max_length=255)


class Score(QuestionBase):
    type: Literal["score"]
    # SDK 0.7.1 permits the degenerate one-level case.
    criteria: list[Content] = Field(min_length=1, max_length=10)


class NoulCriteria(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    true: Content | None = None
    false: Content | None = None


class Noul(QuestionBase):
    type: Literal["noul"]
    criteria: NoulCriteria | None = None


Question = Annotated[Choice | Score | Noul, Field(discriminator="type")]


class SystemOneRequest(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    state: Content
    model: str
    questions: dict[str, Question] = Field(min_length=1, max_length=128)

    @model_validator(mode="after")
    def finite_json(self):
        text = json.dumps(self.model_dump(), ensure_ascii=False, allow_nan=False)
        if len(text.encode()) > 1_000_000:
            raise ValueError("request exceeds 1 MB")
        return self


def render(value):
    return json.dumps(value, ensure_ascii=False, sort_keys=True, allow_nan=False)


def options(question):
    """IDs never enter this function; choice order is canonical, score order is semantic."""
    q = question.model_dump() if isinstance(question, BaseModel) else question
    if q["type"] == "choice":
        return [
            (k, render({"label": k, "description": v})) for k, v in sorted(q["criteria"].items())
        ]
    if q["type"] == "score":
        return [(str(i), render(v)) for i, v in enumerate(q["criteria"])]
    c = q.get("criteria") or {}
    return [
        ("false", render({"answer": "否 / false", "description": c.get("false")})),
        ("true", render({"answer": "是 / true", "description": c.get("true")})),
    ]


def answer(question, probabilities):
    keys = [k for k, _ in options(question)]
    p = [float(x) for x in probabilities]
    if len(p) != len(keys) or any(not math.isfinite(x) or x < 0 for x in p):
        raise ValueError("invalid model probabilities")
    total = sum(p)
    if total <= 0:
        raise ValueError("empty probability mass")
    p = [x / total for x in p]
    kind = question["type"]
    if kind == "noul":
        return {"type": kind, "noul": p[1]}
    entropy = -sum(x * math.log(x) for x in p if x > 0)
    confidence = 1.0 if len(p) == 1 else max(0.0, min(1.0, 1 - entropy / math.log(len(p))))
    result = {
        "type": kind,
        "probabilities": dict(zip(keys, p, strict=True)),
        "confidence": confidence,
    }
    if kind == "choice":
        result["choice"] = keys[max(range(len(p)), key=p.__getitem__)]
    else:
        result["score"] = sum(i * value for i, value in enumerate(p))
        result["legend"] = {str(i): v for i, v in enumerate(question["criteria"])}
    return result