#!/usr/bin/env python3
"""patch_fixed_chunks.py -- POSITION-FIXED prompt chunking when the prefix cache is on (audit 258).

ROOT CAUSE (proven by the audit-257 provenance dump): `chunks(L0)` is
    n = ceil(L0/rows); step = ceil(L0/n)
so `step` is a function of the PROMPT LENGTH. Two requests whose prompts differ in length therefore compute the
SAME prefix rows with DIFFERENT chunk boundaries, i.e. a different row count R per chunk, and the kernels' tiling
is a function of R. The KV that a resumed request inherits was written by the WARM request under the warmer
request's layout, so it is not the KV a cold run of the current prompt would produce.

Evidence: the passing case had len(W) == len(P) (same step -> identical KV, verified per row); the failing cases
had len(W) != len(P), and the dump showed layer 0 identical, layers 39/77 differing from position ~8 -- a
numerical divergence inside the compute, not a bookkeeping error.

FIX: when TF_GLM53_FIXED_CHUNKS is on, chunks are a FIXED grid of `prompt_rows` tokens plus a final remainder.
Every chunk boundary is then a multiple of `prompt_rows`, independent of L0, so a prefix's rows are always
produced by the same tiling and a resumed prefill is bit-exact by construction. Reuse offsets are quantized to
those boundaries, so reused rows always sit in FULL chunks.

Default OFF: with fixed chunks off, the cold path is byte-for-byte what it is today (banked sha unchanged).
usage: patch_fixed_chunks.py <runner.py> [--apply]
"""
import ast, sys
P = sys.argv[1]; apply = "--apply" in sys.argv
src = open(P).read()

old1 = """    def chunks(self, L0: int) -> list[tuple[int, int]]:"""
new1 = """    def chunks(self, L0: int) -> list[tuple[int, int]]:
        if FIXED_CHUNKS:
            # audit 258: a POSITION-FIXED grid, so a prompt prefix is always computed with the same tiling
            # whatever the total length. Required for a bit-exact prefix cache; see the module patch note.
            step = self.prompt_rows
            return [(a, min(a + step, L0)) for a in range(0, L0, step)]"""

old2 = 'KVDIR = os.environ.get("TF_GLM53_KVDIR", "")   # audit 257: per-row KV provenance dump dir'
new2 = (old2 + '\nFIXED_CHUNKS = os.environ.get("TF_GLM53_FIXED_CHUNKS", "0") != "0"   # audit 258: position-fixed grid')

for label, o in (("chunks: fixed grid", old1), ("module: FIXED_CHUNKS", old2)):
    if src.count(o) != 1:
        print("REFUSED: %s: %d occurrences" % (label, src.count(o))); sys.exit(1)
out = src.replace(old2, new2).replace(old1, new1)
ast.parse(out)
print("2 anchors unique; parses OK")
if not apply: print("dry run -- pass --apply"); sys.exit(0)
open(P, "w").write(out); print("APPLIED")