"""Benchmark stop rules, applied identically to AR and speculative decoding.

Same criteria as scripts/bench_speed.StopRule (chantiers/06-freins-vitesse.md, T3): GSM8K
stops after the '#### <answer>' line or at the next few-shot question; HumanEval at the
usual function-end sequences.
"""

from __future__ import annotations

import re


class StopRule:
    SEQS = {"gsm8k": ("\n\nQuestion:",), "humaneval": ("\ndef ", "\nclass ", "\nif __name__", "\nprint(", "\n#")}

    def __init__(self, tok, bench):
        self.tok, self.bench, self.text = tok, bench, ""

    def __call__(self, t):
        self.text += self.tok.decode([t])
        tail = self.text[-48:]
        if self.bench == "gsm8k" and re.search(r"####[^\n]*\d[^\n]*\n", tail):
            return True
        return any(s in tail for s in self.SEQS.get(self.bench, ()))


def make_stop(tok, kind: str | None) -> StopRule | None:
    """A fresh stop rule (they are stateful: one per generation), or None."""
    if kind is None:
        return None
    if kind not in StopRule.SEQS:
        raise ValueError(f"unknown stop rule {kind!r}; available: {sorted(StopRule.SEQS)} or None")
    return StopRule(tok, kind)