"""Prompt and canvas construction. Token spans are exact because every piece is tokenised separately."""
from __future__ import annotations

import random
from dataclasses import dataclass

from .schema import ALIASES, Example

# Bump when anything below changes the token ids, option spans or canvas a trained head was fitted on; every
# feature cache must then be re-extracted and the head retrained. tests/test_format_pin.py pins version 1.
PROMPT_FORMAT_VERSION = 1
SENTINEL = "<<HEV_USER_TEXT>>"
HEADER = "State:\n"
QUESTIONS_HEADER = "\nQuestions:\n"
FOOTER = "\nAnswer as JSON with one option letter per question.\n"
MIN_STATE_TOKENS = 16


@dataclass
class Piece:
    text: str
    kind: str  # header | state | question | option | footer
    qidx: int = -1
    oidx: int = -1
    alias_idx: int = -1


@dataclass
class TokenizedPrompt:
    ids: list[int]
    option_spans: list[list[tuple[int, int]]]
    option_alias: list[list[int]]
    truncated: bool


def render_pieces(example: Example, option_order: dict[int, list[int]] | None = None) -> list[Piece]:
    pieces = [Piece(HEADER, "header"), Piece(example.state.strip() + "\n", "state"),
              Piece(QUESTIONS_HEADER, "header")]
    for qi, q in enumerate(example.questions):
        pieces.append(Piece(f"{q.qid} ({q.type}): {q.instructions}\n", "question", qi))
        order = (option_order or {}).get(qi) or list(range(q.k))
        for pos, oi in enumerate(order):
            pieces.append(Piece(f"  {ALIASES[pos]}: {q.descriptions[oi]}\n", "option", qi, oi, pos))
    pieces.append(Piece(FOOTER, "footer"))
    return pieces


def chat_affixes(tok) -> tuple[list[int], list[int]]:
    text = tok.apply_chat_template([{"role": "user", "content": SENTINEL}],
                                   tokenize=False, add_generation_prompt=True)
    pre, post = text.split(SENTINEL)
    return tok.encode(pre, add_special_tokens=False), tok.encode(post, add_special_tokens=False)


def tokenize_prompt(example: Example, tok, max_tokens: int = 1024,
                    option_order: dict[int, list[int]] | None = None) -> TokenizedPrompt:
    pre, post = chat_affixes(tok)
    pieces = render_pieces(example, option_order)
    enc = [tok.encode(p.text, add_special_tokens=False) for p in pieces]
    budget = max_tokens - len(pre) - len(post)
    fixed = sum(len(e) for e, p in zip(enc, pieces) if p.kind != "state")
    truncated = False
    for i, p in enumerate(pieces):
        if p.kind == "state" and fixed + len(enc[i]) > budget:
            keep = max(budget - fixed, MIN_STATE_TOKENS)
            half = keep // 2
            enc[i] = enc[i][:half] + enc[i][-(keep - half):]
            truncated = True
    ids = list(pre)
    spans = [[(0, 0)] * q.k for q in example.questions]
    alias = [[0] * q.k for q in example.questions]
    for p, e in zip(pieces, enc):
        start = len(ids)
        ids.extend(e)
        if p.kind == "option":
            spans[p.qidx][p.oidx] = (start, len(ids))
            alias[p.qidx][p.oidx] = p.alias_idx
    ids.extend(post)
    if len(ids) > max_tokens:
        raise ValueError(f"{example.id}: prompt has {len(ids)} tokens even after truncation (max {max_tokens})")
    return TokenizedPrompt(ids=ids, option_spans=spans, option_alias=alias, truncated=truncated)


def canvas_ids(example: Example, tok, slot_token_id: int, canvas_length: int,
               eos_id: int = 1) -> tuple[list[int], list[int]]:
    ids: list[int] = []
    slots: list[int] = []
    for qi, q in enumerate(example.questions):
        prefix = '{"' if qi == 0 else ', "'
        ids.extend(tok.encode(prefix + q.qid + '": "', add_special_tokens=False))
        slots.append(len(ids))
        ids.append(slot_token_id)
        ids.extend(tok.encode('"', add_special_tokens=False))
    ids.extend(tok.encode("}", add_special_tokens=False))
    if len(ids) > canvas_length:
        raise ValueError(f"{example.id}: canvas needs {len(ids)} tokens, canvas_length is {canvas_length}")
    ids.extend([eos_id] * (canvas_length - len(ids)))
    return ids, slots


def alias_token_ids(tok, n: int) -> list[int]:
    out = []
    for a in ALIASES[:n]:
        ids = tok.encode(a, add_special_tokens=False)
        if len(ids) != 1:
            raise ValueError(f"alias {a!r} is not a single token: {ids}")
        out.append(ids[0])
    return out


def shuffled_order(example: Example, rng: random.Random) -> dict[int, list[int]]:
    order = {}
    for qi, q in enumerate(example.questions):
        if q.type == "choice":
            perm = list(range(q.k))
            rng.shuffle(perm)
            order[qi] = perm
    return order