#!/usr/bin/env python3
"""patch_prefill_offset.py -- INCREMENT 1 of the full prefix cache (audit 252/253).

Adds an `offset` to the prompt path so a request can prefill ONLY a suffix, starting at absolute position
`offset`, while the KV for positions [0, offset) is already valid in the shared persistent `fused.State`.

`offset=0` (the default) must be BYTE-IDENTICAL to today's behaviour -- that is the acceptance test.

Anchors are exact and unique; dry-run unless --apply.
usage: patch_prefill_offset.py <runner.py> [--apply]
"""
import ast, sys

P = sys.argv[1]
apply = "--apply" in sys.argv
src = open(P).read()

# 1. prefill(): accept and forward the offset (both the profiled and the plain branch).
old1 = """    def prefill(self, prompt: list[int]) -> torch.Tensor:
        flag = PROFILE_FLAG + "_PREFILL"
        if not os.path.exists(flag):
            return self._prefill(prompt)"""
new1 = """    def prefill(self, prompt: list[int], offset: int = 0) -> torch.Tensor:
        flag = PROFILE_FLAG + "_PREFILL"
        if not os.path.exists(flag):
            return self._prefill(prompt, offset)"""

old2 = """            out = self._prefill(prompt)
            torch.cuda.synchronize()"""
new2 = """            out = self._prefill(prompt, offset)
            torch.cuda.synchronize()"""

# 2. _prefill(): chunk over the ABSOLUTE range [offset, L0). offset=0 reduces to today's loop exactly.
old3 = """    def _prefill(self, prompt: list[int]) -> torch.Tensor:
        \"\"\"Every prompt row through the target (and the MTP layer); returns the last row's logits [1, vocab].\"\"\"
        w = self.w
        L0 = len(prompt)
        toks = torch.tensor(prompt, dtype=torch.long, device=w.device)
        logits = None
        if self.drafter is not None:
            self.drafter.reset()
        if self.cap is not None:
            self.cap.reset()
        for a, e in self.chunks(L0):
            lg = self.prefill_chunk(self.st, toks, a, e, L0, taps=True)
            if lg is not None:
                logits = lg
        return logits"""
new3 = """    def _prefill(self, prompt: list[int], offset: int = 0) -> torch.Tensor:
        \"\"\"Rows ``offset``..L0-1 of ``prompt`` through the target (and the MTP layer); returns the last row's logits.

        ``offset`` > 0 (audit 253): the KV for positions [0, offset) is already valid in this Runner's shared
        ``State`` (it was written by an earlier request whose prompt is a prefix of this one), so only the suffix
        is prefilled. Positions stay ABSOLUTE -- ``prefill_chunk`` already sets ``st.pos`` from ``a`` and derives
        the index-key range ``T`` from ``e``, so nothing else changes. ``offset=0`` replays the old loop exactly.
        \"\"\"
        w = self.w
        L0 = len(prompt)
        toks = torch.tensor(prompt, dtype=torch.long, device=w.device)
        logits = None
        if self.drafter is not None:
            self.drafter.reset()
        if self.cap is not None:
            self.cap.reset()
        for a0, e0 in self.chunks(L0 - offset):
            a, e = a0 + offset, e0 + offset
            lg = self.prefill_chunk(self.st, toks, a, e, L0, taps=True)
            if lg is not None:
                logits = lg
        return logits"""

# 3. generate(): forward the offset to prefill. `P = L0` and `sample(lg, L0)` are already correct, because the
#    resumed suffix still ends at absolute position L0.
old4 = """        L0 = len(prompt)
        lg = self.prefill(prompt)"""
new4 = """        L0 = len(prompt)
        lg = self.prefill(prompt, offset)"""
old5 = """    def generate(self, prompt: list[int], max_tokens: int, sample: Callable, stop: Callable[[int], bool],
                 on_tokens: Callable[[list[int]], Any], k: int, mode: str | None = None,
                 sampling=None) -> dict[str, Any]:"""
new5 = """    def generate(self, prompt: list[int], max_tokens: int, sample: Callable, stop: Callable[[int], bool],
                 on_tokens: Callable[[list[int]], Any], k: int, mode: str | None = None,
                 sampling=None, offset: int = 0) -> dict[str, Any]:"""

for label, t, old in (("prefill: forward offset (plain)", src, old1),
                      ("prefill: forward offset (profiled)", src, old2),
                      ("_prefill: absolute chunk range", src, old3),
                      ("generate: pass offset to prefill", src, old4),
                      ("generate: offset parameter", src, old5)):
    if t.count(old) != 1:
        print("REFUSED: %s: found %d occurrences (need exactly 1)" % (label, t.count(old)))
        sys.exit(1)

out = (src.replace(old5, new5).replace(old1, new1).replace(old2, new2)
          .replace(old3, new3).replace(old4, new4))
ast.parse(out)
print("all 5 anchors unique; parses OK")
if not apply:
    print("dry run -- pass --apply to write")
    sys.exit(0)
open(P, "w").write(out)
print("APPLIED")