#!/usr/bin/env python3
import ast, sys
PC, RUN, ENG = sys.argv[1], sys.argv[2], sys.argv[3]
apply = "--apply" in sys.argv
pc, rn, en = open(PC).read(), open(RUN).read(), open(ENG).read()
OLD1 = '    def __init__(self) -> None:\n        self.ids: list = []\n        self.enabled = TF_PREFIX_CACHE'
NEW1 = '    def __init__(self) -> None:\n        self.ids: list = []\n        self.canonical_end: int = 0        # audit 260: highest row a CANONICAL PREFILL chunk wrote\n        self.enabled = TF_PREFIX_CACHE'
OLD2 = '        if not self.enabled or not self.ids:\n            return 0\n        n = common_prefix(self.ids, prompt)\n        if n < KEEP_MIN or n >= len(prompt):\n            return 0\n        return n'
NEW2 = "        if not self.enabled or not self.ids:\n            return 0\n        n = common_prefix(self.ids, prompt)\n        # audit 260: never offer a row decode wrote. canonical_end is the previous request's prompt\n        # length -- the highest position a full prefill chunk produced.\n        n = min(n, self.canonical_end, len(prompt) - 1)\n        if n < KEEP_MIN:\n            return 0\n        return n"
OLD3 = '    def record(self, ids: list) -> None:\n        """The State now holds ``ids`` (call after a request, with the full prompt + committed reply)."""\n        if self.enabled:\n            self.ids = list(ids)'
NEW3 = '    def record(self, ids: list, canonical_end: int) -> None:\n        """Publish the record. canonical_end = this request PROMPT length, the highest row a canonical\n        prefill chunk wrote. Generated tokens past it are decode-written and must never be reused."""\n        if self.enabled:\n            self.ids = list(ids)\n            self.canonical_end = max(0, int(canonical_end))'
OLD4 = 'FIXED_CHUNKS = os.environ.get("TF_GLM53_FIXED_CHUNKS",\n                              os.environ.get("TF_GLM53_PREFIX_CACHE", "1")) != "0"   # audit 259: bound to the cache'
NEW4 = '# audit 260: a HARD guard, not just a default. Fixed chunking is what makes reuse bit-exact (audit 258),\n# so the cache being on forces it on; an explicit TF_GLM53_FIXED_CHUNKS=0 cannot restore the old failure.\nPREFIX_ON = os.environ.get("TF_GLM53_PREFIX_CACHE", "1") != "0"\nFIXED_CHUNKS = PREFIX_ON or (os.environ.get("TF_GLM53_FIXED_CHUNKS", "0") != "0")'
OLD5 = '            self.prefix.record([*prompt, *out[:-1]])'
NEW5 = '            self.prefix.record([*prompt, *out[:-1]], canonical_end=len(prompt))'
for lbl, t, o in ((1,pc,OLD1),(2,pc,OLD2),(3,pc,OLD3),(4,rn,OLD4),(5,en,OLD5)):
    if t.count(o) != 1:
        print("REFUSED anchor", lbl, t.count(o)); sys.exit(1)
pc2 = pc.replace(OLD1,NEW1).replace(OLD2,NEW2).replace(OLD3,NEW3)
rn2 = rn.replace(OLD4,NEW4)
en2 = en.replace(OLD5,NEW5)
ast.parse(pc2); ast.parse(rn2); ast.parse(en2)
print("5 anchors unique; all parse OK")
if not apply: print("dry run -- pass --apply"); sys.exit(0)
open(PC,"w").write(pc2); open(RUN,"w").write(rn2); open(ENG,"w").write(en2)
print("APPLIED")