"""Prompt rendering. Training and inference share this exact format."""

from __future__ import annotations

from .schema import Question, normalize_state

SPECIAL_TOKENS = [
    "<|state|>",
    "<|question|>",
    "<|type|>",
    "<|options|>",
    "<|answer|>",
    "<|opt|>",
]

OPTION_MARKER = "<|opt|>"


def option_texts(q: Question) -> list[str]:
    if q.kind == "noul":
        return ["false", "true"]
    if q.kind == "choice":
        out = []
        for key, desc in zip(q.keys, q.descriptions):
            out.append(key if desc == key else f"{key} - {desc}")
        return out
    return list(q.descriptions)


def render_prompt(state: str, q: Question) -> str:
    state = normalize_state(state)
    if q.kind == "score":
        instruction = q.instructions + " Options are ordered from lowest to highest."
    else:
        instruction = q.instructions
    lines = "\n".join(f"{i}. {t} {OPTION_MARKER}" for i, t in enumerate(option_texts(q)))
    return (
        f"<|state|>\n{state}\n"
        f"<|question|>\n{instruction}\n"
        f"<|type|>\n{q.kind}\n"
        f"<|options|>\n{lines}\n"
        f"<|answer|>"
    )