"""Render a decision as model input, with one readout per question.

Options are listed under **neutral answer codes**, and the answer is read at a marker after the
options from the LM-head rows of those codes:

```
<state>
…state text, JSON or images…
</state>
<question type=choice>
Which team should handle this ticket?
BK = billing: Payments and refunds
QX = tech: Technical errors
Answer:⟨readout⟩
```

* **Codes.** Every question draws its codes from a pool of single-token strings, so no code
  carries meaning across decisions.
* **Noul** (yes/no) is rendered as two options, ``holds`` and ``does_not_hold``, each with its
  criterion, in random order under neutral codes.
* **Score** levels keep their order (codes are still neutral); the level index is written out.
* **Isolated** mode (default) gives each question its own sequence over the shared state, so
  answers do not depend on which other questions were asked. **Joint** mode puts every question
  in one sequence with one readout each.
* **Images** open the state block, in order, as the ids the image encoder returns
  (``<|vision_start|>``, one ``<|image_pad|>`` per merged patch, ``<|vision_end|>``).

The readout index is the position of the last token of the ``Answer:`` marker.
"""

from __future__ import annotations

import json
import random
from dataclasses import dataclass, field
from typing import Any, Callable, Protocol, Sequence

from glance_runtime.schema import Decision, Question

NOUL_HOLDS, NOUL_NOT = "holds", "does_not_hold"
ANSWER_MARKER = "\nAnswer:"


class ImagesUnsupported(ValueError):
    """The decision has images and the model was given no image encoder."""


class Tokenizer(Protocol):
    def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]:
        ...


@dataclass
class Readout:
    question: str
    type: str
    position: int  # token index of the readout (last token of the answer marker)
    codes: list[str]  # code shown for each option, in display order
    labels: list[str]  # the question's labels, in display order
    order: list[int]  # display index -> index into the question's labels


@dataclass
class Example:
    input_ids: list[int]
    readouts: list[Readout]
    text: str
    decision_id: str = ""
    vision: list[Any] = field(default_factory=list)  # one image-encoder payload per image, in order


def _state_text(d: Decision) -> str:
    return d.state if isinstance(d.state, str) else json.dumps(d.state, ensure_ascii=False, indent=1)


def _options(q: Question, rng: random.Random) -> tuple[list[str], list[str | None], list[int]]:
    """Display names, definitions and order (display index -> label index)."""
    n = len(q.labels)
    defs = list(q.definitions) if q.definitions else [None] * n
    if q.type == "noul":
        holds = defs[1] or "the statement in the question is true"
        not_holds = defs[0] or "the statement in the question is false"
        order = [1, 0] if rng.random() < 0.5 else [0, 1]
        names = {1: NOUL_HOLDS, 0: NOUL_NOT}
        texts = {1: holds, 0: not_holds}
        return [names[i] for i in order], [texts[i] for i in order], order
    order = list(range(n))  # choice options keep the caller's order; score keeps its order
    if q.type == "score":
        return [f"level {i} of {n - 1}" for i in order], [q.labels[i] for i in order], order
    return [q.labels[i] for i in order], [defs[i] for i in order], order


def question_block(q: Question, codes: list[str], rng: random.Random) -> tuple[str, list[int]]:
    names, defs, order = _options(q, rng)
    lines = [f"<question type={q.type}>", q.instructions.strip() or "Choose the option that fits."]
    for code, nm, df in zip(codes, names, defs):
        lines.append(f"{code} = {nm}: {df}" if df else f"{code} = {nm}")
    return "\n".join(lines) + "\n</question>", order


def render(d: Decision, tokenizer: Tokenizer, code_pool: Sequence[str], rng: random.Random, *,
           joint: bool = False, max_tokens: int | None = None,
           image_encoder: Callable[[dict[str, Any]], tuple[list[int], Any]] | None = None) -> list[Example]:
    """Render ``d`` into one example per question (isolated) or one example (joint)."""
    if d.images and image_encoder is None:
        raise ImagesUnsupported(f"{len(d.images)} image(s) and no image encoder")
    encoded = [image_encoder(im) for im in d.images] if d.images else []  # once per decision, shared by its examples
    state = f"<state>\n{_state_text(d)}\n</state>"
    blocks = []
    for name, q in d.questions.items():
        n = len(q.labels)
        if n > len(code_pool):
            raise ValueError(f"{name}: {n} options but only {len(code_pool)} codes")
        codes = rng.sample(list(code_pool), n)
        text, order = question_block(q, codes, rng)
        blocks.append((text, Readout(name, q.type, -1, codes, [q.labels[i] for i in order], order)))

    groups = [blocks] if joint else [[b] for b in blocks]
    out = []
    for group in groups:
        ids: list[int] = []
        text_parts: list[str] = []
        readouts = []

        def add(s: str) -> None:
            ids.extend(tokenizer.encode(s, add_special_tokens=False))
            text_parts.append(s)

        def add_state(suffix: str = "") -> None:
            if not encoded:
                add(state + suffix)
                return
            add("<state>\n")
            for image_ids, _ in encoded:
                ids.extend(image_ids)
                text_parts.append("<image>")
                add("\n")
            add(f"{_state_text(d)}\n</state>{suffix}")

        add_state("\n")
        for qtext, ro in group:
            add(qtext)
            add(ANSWER_MARKER)
            ro.position = len(ids) - 1
            readouts.append(ro)
            add("\n")
        if max_tokens is not None and len(ids) > max_tokens:
            raise ValueError(f"{len(ids)} tokens exceeds {max_tokens}")
        out.append(Example(ids, readouts, "".join(text_parts), d.id, [v for _, v in encoded]))
    return out