"""
runner_patch.py — swap HyperExpertRunner onto FastGenerator WITHOUT editing
runtime_adapters.py. Reversible, A/B-able.

    from runner_patch import install, uninstall
    install(runner)      # runner._generate now uses cached decode (sliding)
    uninstall(runner)    # back to stock

GATED: only install after divergence_quality.py says benign AND
help/shatter + smoke re-run matches native within noise.

Preserves every stock control byte-for-byte (verified identical output with
rep_pen=1.3 + no_repeat_3gram + pad masking on a real runner instance): the
loop body is stock _generate's; only the logits source changes (full forward
per step -> FastGenerator prefill/step). Deltas set before the call are
picked up automatically (snapshot-at-prefill).

One stated semantic difference: stock _generate rolls a sliding window when
len(prompt)+new exceeds block_size; a cache cannot roll, so we pre-crop the
prompt to (block_size - max_new). Your harness already budgets context so
this should never trigger; if it does, the crop is logged.
"""
import types
from time import perf_counter

from fast_generate import FastGenerator


def _generate_fast(self, prompt_text: str, max_new_tokens: int) -> str:
    torch = self.torch
    g = self.gen
    from runtime_adapters import _banned_ngram_tokens, log
    ids = self.tok.encode(prompt_text)[: self.block_size]
    if len(ids) + max_new_tokens > self.block_size:
        log(f"fast decode: pre-cropping prompt {len(ids)} -> "
            f"{self.block_size - max_new_tokens} (cache cannot roll)",
            tag="GENERATE")
        ids = ids[-(self.block_size - max_new_tokens):]
    idx = torch.tensor([ids], dtype=torch.long, device=self.device)
    t0 = perf_counter()
    fg = FastGenerator(self.adapted.model)          # sliding; sees live deltas
    logits_next = fg.prefill(idx)
    new_ids: list[int] = []
    n_blocked = 0
    for _ in range(max_new_tokens):
        logits = logits_next.float()                # (1, V)
        if g.mask_pad_vocab and self.tok_vocab and self.tok_vocab < logits.size(-1):
            logits[:, self.tok_vocab:] = -float("inf")
        if g.repetition_penalty and g.repetition_penalty != 1.0:
            seen = torch.unique(idx[0])
            vals = logits[0, seen]
            logits[0, seen] = torch.where(vals > 0, vals / g.repetition_penalty,
                                          vals * g.repetition_penalty)
        if g.no_repeat_ngram_size and g.no_repeat_ngram_size >= 2:
            banned = _banned_ngram_tokens(idx[0].tolist(), g.no_repeat_ngram_size)
            if banned:
                n_blocked += len(banned)
                logits[0, list(banned)] = -float("inf")
        if g.temperature <= 0:
            nxt = int(logits.argmax(-1))
        else:
            lg = logits / g.temperature
            if g.top_k:
                v, _ = torch.topk(lg, min(g.top_k, lg.size(-1)))
                lg[lg < v[:, [-1]]] = -float("inf")
            nxt = int(torch.multinomial(torch.softmax(lg, -1), 1))
        new_ids.append(nxt)
        idx = torch.cat([idx, torch.tensor([[nxt]], device=self.device)], dim=1)
        if g.stop_on_eos and nxt in self.stop_ids:
            new_ids.pop()
            break
        logits_next = fg.step(torch.tensor([[nxt]], device=self.device))
    extra = (f", rep_pen={g.repetition_penalty}" if g.repetition_penalty != 1.0 else "")
    extra += (f", no_repeat_{g.no_repeat_ngram_size}gram(blocked {n_blocked})"
              if g.no_repeat_ngram_size >= 2 else "")
    log(f"generate[FAST]: prompt {len(ids)} tok -> +{len(new_ids)} tok "
        f"in {perf_counter()-t0:.2f}s "
        f"({'greedy' if g.temperature <= 0 else f'T={g.temperature}'}{extra})",
        tag="GENERATE")
    return self.tok.decode(new_ids)


def install(runner):
    if not hasattr(runner, "_generate_stock"):
        runner._generate_stock = runner._generate
    runner._generate = types.MethodType(_generate_fast, runner)
    return runner


def uninstall(runner):
    if hasattr(runner, "_generate_stock"):
        runner._generate = runner._generate_stock
    return runner