#!/usr/bin/env python3
"""patch_offset_align2.py -- CORRECTED increment 4: run WHOLE cold chunks, skip whole chunks (audit 255).

The first attempt (audit 255a) quantized the offset to a cold boundary but then still called
`prefill_chunk(..., max(a, offset), e, ...)`, i.e. it started the chunk at the OFFSET. The last chunk's row count
R = e - offset then differed from a cold run's R = e - a, and because the kernels' tiling is a function of R the
reduction order changed -- the same defect, just moved. Measured: unchanged 2/6 and 4/6 mismatches.

CORRECT RULE: run the COLD layout's chunks unchanged -- full `(a, e)` ranges -- and skip only chunks that lie
entirely at or below the offset. Every executed chunk then has exactly the R a cold run gives it, so the KV rows
are produced identically. The resume still skips every whole chunk below the offset (the bulk of the prompt).

`offset` is aligned once in `generate` to the largest cold chunk boundary <= the request's reuse length, so the
value reported as `cached`, passed to the drafter, and used to decide skipped chunks are all the same number.

usage: patch_offset_align2.py <runner.py> [--apply]
"""
import ast, sys
P = sys.argv[1]; apply = "--apply" in sys.argv
src = open(P).read()

# (a) replace the first attempt's loop with whole-chunk execution
old1 = """        ch = self.chunks(L0)                 # the COLD layout for this prompt length
        if offset:
            # audit 255: the chunk boundaries are a function of L0, so a resumed run must reuse the COLD layout
            # exactly. Quantize the offset down to the nearest cold boundary; re-prefilling a partial chunk
            # (<= one chunk) costs little and keeps R -- and therefore the kernels' tiling -- identical to a
            # cold run. Without this, the same KV row is produced by a different reduction order and a near-tie
            # can flip under greedy decoding.
            offset = max((a for a, _e in ch if a <= offset), default=0)
        for a, e in ch:
            if e <= offset:
                continue                     # whole chunk already in the state
            lg = self.prefill_chunk(self.st, toks, max(a, offset), e, L0, taps=True)
            if lg is not None:
                logits = lg
        return logits"""
new1 = """        for a, e in self.chunks(L0):         # audit 255: the COLD layout, whole chunks only
            if e <= offset:
                continue                     # this whole chunk is already valid in the state
            lg = self.prefill_chunk(self.st, toks, a, e, L0, taps=True)
            if lg is not None:
                logits = lg
        return logits"""

# (b) the alignment helper
old2 = """    @staticmethod
    def _check_cache_fits(w, capacity: int, slots: int) -> None:"""
new2 = """    def align_offset(self, L0: int, offset: int) -> int:
        \"\"\"The largest COLD chunk boundary <= ``offset`` (0 when none) -- audit 255.

        A resumed run must execute only whole chunks of the cold layout: the kernels' tiling is a function of a
        chunk's row count R, so starting one chunk at the offset (rather than at its own boundary) changes R and
        therefore the reduction order, which is enough to flip a near-tie under greedy decoding. Quantizing the
        offset down to a boundary costs at most one chunk of re-prefill and makes the reused rows bit-identical.
        \"\"\"
        if not offset:
            return 0
        return max((a for a, _e in self.chunks(L0) if a <= offset), default=0)

    @staticmethod
    def _check_cache_fits(w, capacity: int, slots: int) -> None:"""

# (c) align once in generate, so prefill, the drafter and the reported `cached` all agree
old3 = """        L0 = len(prompt)
        lg = self.prefill(prompt, offset)"""
new3 = """        L0 = len(prompt)
        offset = self.align_offset(L0, offset)      # audit 255: whole cold chunks only
        lg = self.prefill(prompt, offset)"""

for label, old in (("_prefill: whole chunks", old1), ("Runner.align_offset", old2),
                   ("generate: align once", old3)):
    if src.count(old) != 1:
        print("REFUSED: %s: %d occurrences (need 1)" % (label, src.count(old))); sys.exit(1)

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