"""A fast, token-identical System One prompt encoder (qev_fast.systemone.encode without decider's per-piece
transformers calls).

decider.prompt.build (state-first, plain layout, no shuffling, max_options = MAX_OPTIONS) encodes, per request:
- "Context:\\n" + the rendered state, as one string;
- per question row k, either one string (<= NARROW options: header, "\\n(A) option" lines, answer tail), or (wide) the
  header, then per option the ids of "\\n(" + its label token + the ids of ") option", then the answer tail.
FastEncoder builds exactly those strings (decider's own render_state / render_question / plan_rows and
systemone.neutralize_options), encodes all strings of all requests of a call in one Rust `encode_batch` of the
tokenizer backend (the same tokenizer object transformers' fast tokenizer wraps, so the same ids), caches the strings
that repeat (option texts, whole narrow question blocks of fixed schemas) and assembles ids / slots / nopts / golds.
Anything it does not model (more than MAX_OPTIONS options) goes to the reference encoder; `stats` counts both.
"""

import json

from decider.prompt import LETTERS, MAX_OPTIONS, NARROW, label_table
from decider.systemone import plan_rows, render_question, render_state, row_types

from . import systemone as S1


class FastEncoder:
    def __init__(self, tok, cache_size: int = 200_000):
        self.tok = tok
        self.raw = tok.backend_tokenizer  # tokenizers.Tokenizer (Rust)
        _, self.lab_ids, self.open_ids = label_table(tok)
        self.cache: dict = {}
        self.cache_size = cache_size
        self.stats = {"fast": 0, "fallback": 0, "cache_hits": 0, "strings_encoded": 0}

    # ---- planning: the strings decider.prompt.build encodes ---------------------------------------------------------
    def _plan(self, state, questions: dict):
        ctx = render_state(state)  # as S1.encode (callers pass a string or what render_state accepts)
        rqs = {k: render_question(v) for k, v in questions.items()}
        flat, index = plan_rows(rqs, False)
        multi = len(flat) > 1
        parts = [("s", "Context:\n" + ctx, False)]  # (kind, string, cacheable)
        rows = []
        for k, r in enumerate(flat):
            opts = S1.neutralize_options(r["options"])
            if len(opts) > MAX_OPTIONS:
                return None  # option sampling (decider keeps the gold + abstain options): the reference handles it
            head = f"\n\nQuestion{' ' + str(k + 1) if multi else ''}: {r['question']}\nOptions:"
            tail = f"\nAnswer{' ' + str(k + 1) if multi else ''}: ("
            if len(opts) <= NARROW:
                parts.append(("s", head + "".join(f"\n({LETTERS[j]}) {o}" for j, o in enumerate(opts)) + tail, True))
                rows.append(("n", len(parts) - 1, len(opts)))
            else:
                h = len(parts)
                parts.append(("s", head, True))
                for o in opts:
                    parts.append(("s", f") {o}", True))
                parts.append(("s", tail, True))
                rows.append(("w", h, len(opts)))
        return rqs, index, parts, rows

    # ---- encoding -----------------------------------------------------------------------------------------------
    def encode_many(self, requests: list) -> list:
        """[(state, questions)] -> per request S1.encode's dict (+ "golds"), or an exception instance for a request
        the prompt builder refuses (ValueError / KeyError / TypeError / AssertionError, as S1.encode raises)."""
        plans, todo = [], {}
        for state, qs in requests:
            try:
                p = self._plan(state, qs)
            except (ValueError, KeyError, TypeError, AssertionError) as e:
                plans.append(e)
                continue
            plans.append(p if p is not None else ("fallback", state, qs))
            if p is not None:
                for _kind, s, cacheable in p[2]:
                    if not (cacheable and s in self.cache):
                        todo[s] = None
                    else:
                        self.stats["cache_hits"] += 1
        if todo:
            strs = list(todo)
            enc = self.raw.encode_batch(strs, add_special_tokens=False)
            for s, e in zip(strs, enc, strict=True):
                todo[s] = e.ids
            self.stats["strings_encoded"] += len(strs)
        out = []
        for p in plans:
            if isinstance(p, Exception):
                out.append(p)
                continue
            if p[0] == "fallback":
                self.stats["fallback"] += 1
                try:
                    e = S1.encode(self.tok, p[1], p[2])
                    e["golds"] = None
                    out.append(e)
                except (ValueError, KeyError, TypeError, AssertionError) as ex:
                    out.append(ex)
                continue
            rqs, index, parts, rows = p

            def ids_of(i):
                s, cacheable = parts[i][1], parts[i][2]
                v = todo.get(s)
                if v is None:
                    v = self.cache[s]
                elif cacheable and len(self.cache) < self.cache_size:
                    self.cache[s] = v
                return v

            ids = list(ids_of(0))
            state_len = len(ids)  # "Context:\n" + the state, encoded as one string: the row's state prefix
            slots, nopts = [], []
            for kind, i, n in rows:
                if kind == "n":
                    ids += ids_of(i)
                else:
                    ids += ids_of(i)
                    for j in range(n):
                        ids += self.open_ids
                        ids.append(self.lab_ids[j])
                        ids += ids_of(i + 1 + j)
                    ids += ids_of(i + 1 + n)
                slots.append(len(ids) - 1)
                nopts.append(n)
            self.stats["fast"] += 1
            out.append({"ids": ids, "slots": slots, "nopts": nopts, "golds": [0] * len(rows), "rqs": rqs,
                        "index": index, "types": row_types(rqs, index), "state_len": state_len})
        return out

    def encode(self, state, questions: dict) -> dict:
        e = self.encode_many([(state, questions)])[0]
        if isinstance(e, Exception):
            raise e
        return e