#!/usr/bin/env python3
"""patch_finalize_cache.py -- make the prefix cache production-default and safe (audit 259).

1. `PREFIX_CACHE` defaults ON (the goal: enabled by default).
2. `FIXED_CHUNKS` is BOUND to it: when TF_GLM53_FIXED_CHUNKS is unset it follows TF_GLM53_PREFIX_CACHE. Fixed
   chunking is what makes reuse bit-exact (audit 258), so the cache must never be enabled without it.
   With the cache off, fixed chunks are off too, so the cold path is byte-for-byte unchanged (banked sha).
3. `record()` now publishes only what is provably materialized: `prompt + out[:-1]`. `generate` documents
   `P = L0  # tok sits at P, not yet in the caches`, so claiming the whole `out` overstates the valid frontier by
   at least one token (both planners flagged this independently).

usage: patch_finalize_cache.py <prefix_cache.py> <runner.py> <engine.py> [--apply]
"""
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 = 'TF_PREFIX_CACHE = os.environ.get("TF_GLM53_PREFIX_CACHE", "0") != "0"'
new1 = 'TF_PREFIX_CACHE = os.environ.get("TF_GLM53_PREFIX_CACHE", "1") != "0"   # audit 259: ON by default'

old2 = 'FIXED_CHUNKS = os.environ.get("TF_GLM53_FIXED_CHUNKS", "0") != "0"   # audit 258: position-fixed grid'
new2 = ('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')

old3 = """            self.prefix.record([*prompt, *out])          # audit 253: the State now holds prompt + reply"""
new3 = """            # audit 259: publish only what is provably materialized. `generate` documents
            # `P = L0  # tok sits at P, not yet in the caches`, so the final emitted token's row may not be
            # written; claiming the whole `out` overstates the valid frontier by at least one token.
            self.prefix.record([*prompt, *out[:-1]])"""

for label, t, o in (("prefix_cache: default ON", pc, old1), ("runner: bind FIXED_CHUNKS", rn, old2),
                    ("engine: record out[:-1]", en, old3)):
    if t.count(o) != 1:
        print("REFUSED: %s: %d" % (label, t.count(o))); sys.exit(1)
pc2, rn2, en2 = pc.replace(old1,new1), rn.replace(old2,new2), en.replace(old3,new3)
ast.parse(pc2); ast.parse(rn2); ast.parse(en2)
print("3 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")