"""Shared question and record types.

Public API types (Choice, Score, Noul) mirror Jev/Laya. Internally every question is
normalized to a `Question` with an ordered list of `Option`s so the data pipeline, the
serializer and the model all speak one format. Nothing here depends on torch.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Literal, Union

QType = Literal["choice", "score", "noul"]
QTYPES: tuple[QType, ...] = ("choice", "score", "noul")


# ---------------------------------------------------------------- public API types

@dataclass
class Choice:
    """Pick one of several unordered options. criteria: {key: description} or [key, ...]."""
    instructions: str
    criteria: Union[dict[str, str], list[str]]


@dataclass
class Score:
    """Place the state on an ordered scale. criteria: level descriptions, lowest first (2-10)."""
    instructions: str
    criteria: list[str]


@dataclass
class Noul:
    """Probability that the statement in `instructions` holds for the state."""
    instructions: str


PublicQuestion = Union[Choice, Score, Noul]


# ---------------------------------------------------------------- normalized form

@dataclass
class Option:
    key: str   # short label, e.g. "billing"; for score levels the level index as str
    text: str  # description; may be "" when the key says it all


@dataclass
class Question:
    type: QType
    instructions: str
    options: list[Option] = field(default_factory=list)  # empty for noul

    @property
    def n_options(self) -> int:
        """Number of scored options; noul is scored as two options (false, true)."""
        return 2 if self.type == "noul" else len(self.options)

    @classmethod
    def from_public(cls, q: PublicQuestion) -> "Question":
        if isinstance(q, Choice):
            items = q.criteria.items() if isinstance(q.criteria, dict) else ((k, "") for k in q.criteria)
            return cls("choice", q.instructions, [Option(str(k), str(v)) for k, v in items])
        if isinstance(q, Score):
            return cls("score", q.instructions, [Option(str(i), t) for i, t in enumerate(q.criteria)])
        if isinstance(q, Noul):
            return cls("noul", q.instructions, [])
        raise TypeError(f"not a question type: {type(q).__name__}")

    def validate(self) -> None:
        assert self.type in QTYPES, self.type
        if self.type == "noul":
            assert not self.options, "noul questions have no options"
        elif self.type == "score":
            assert 2 <= len(self.options) <= 10, "score needs 2-10 levels"
        else:
            assert 2 <= len(self.options) <= 255, "choice needs 2-255 options"
            keys = [o.key for o in self.options]
            assert len(set(keys)) == len(keys), "choice keys must be unique"


# ---------------------------------------------------------------- training record

@dataclass
class Record:
    """One (state, question, gold answer) training/eval example = one JSONL line.

    label: index into question.options for choice/score; 0 (false) or 1 (true) for noul.
    """
    id: str
    source: str  # dataset of origin, e.g. "ag_news"
    state: str
    question: Question
    label: int

    def validate(self) -> None:
        self.question.validate()
        assert 0 <= self.label < self.question.n_options, (self.id, self.label)
        assert self.state.strip(), self.id

    def to_json(self) -> str:
        q = self.question
        return json.dumps({
            "id": self.id, "source": self.source, "state": self.state,
            "type": q.type, "instructions": q.instructions,
            "options": [{"key": o.key, "text": o.text} for o in q.options],
            "label": self.label,
        }, ensure_ascii=False)

    @classmethod
    def from_json(cls, line: str) -> "Record":
        d = json.loads(line)
        q = Question(d["type"], d["instructions"], [Option(o["key"], o["text"]) for o in d["options"]])
        return cls(d["id"], d["source"], d["state"], q, int(d["label"]))


def read_jsonl(path: str) -> list[Record]:
    with open(path, encoding="utf-8") as f:
        return [Record.from_json(line) for line in f if line.strip()]


def write_jsonl(path: str, records: list[Record]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for r in records:
            f.write(r.to_json() + "\n")