"""Grammar-constrained LogitsProcessor (Python port of src/logits.js).

At each generation step:
  1. Decode the generated suffix back to text.
  2. Advance the grammar NFA by that text.
  3. For every candidate token id, check whether appending its decoded text
     keeps the NFA alive; mask losers to -inf (via the shared LegalCache).
  4. EOS is allowed only once the NFA has reached an accept state.

Per-token decoding can disagree with BPE sequence-decoding in edge cases
(merged punctuation, etc.); for the acrostic patterns we care about this
approximation is fine.
"""

from __future__ import annotations

import time

from transformers import LogitsProcessor

from masking import LegalCache


def build_token_text_table(tokenizer, vocab_size):
    """One-shot build of tokenId -> text, using per-token decode. Special tokens
    decode to '' under skip_special_tokens=True, which we treat as
    "disallowed" (empty string)."""
    texts = tokenizer.batch_decode(
        [[i] for i in range(vocab_size)], skip_special_tokens=True
    )
    return [t if isinstance(t, str) else "" for t in texts]


class GrammarLogitsProcessor(LogitsProcessor):
    def __init__(self, grammar, tokenizer, token_text, eos_token_ids=(), legal_cache=None):
        super().__init__()
        self.grammar = grammar
        self.tokenizer = tokenizer
        self.token_text = token_text
        self.cache = legal_cache or LegalCache(grammar, token_text, eos_token_ids)
        self.prompt_length = None
        self.stats = _fresh_stats()

    def reset(self):
        self.prompt_length = None
        self.stats = _fresh_stats()

    def __call__(self, input_ids, scores):
        t_entry = time.perf_counter()
        ids = input_ids[0]
        if self.prompt_length is None:
            self.prompt_length = ids.shape[0]

        generated = ids[self.prompt_length:].tolist()
        text = (
            self.tokenizer.decode(generated, skip_special_tokens=True)
            if generated
            else ""
        )

        state = self.grammar.advance(self.grammar.initial, text)
        data = scores[0]

        if state == -1:
            # Already violated; nothing useful to do without rewinding. Let the
            # original logits through so generation at least terminates.
            self._record(time.perf_counter() - t_entry, -1)
            return scores

        illegal = self.cache.illegal_tensor(state)
        data[illegal.to(data.device)] = float("-inf")

        self._record(time.perf_counter() - t_entry, int((~illegal).sum().item()))
        return scores

    def _record(self, dt, survivors):
        st = self.stats
        st["calls"] += 1
        st["total_ms"] += dt * 1000.0
        st["per_step"].append({"ms": dt * 1000.0, "survivors": survivors})


def _fresh_stats():
    return {"calls": 0, "total_ms": 0.0, "per_step": []}