#!/usr/bin/env python3
"""patch_engine_cache.py -- INCREMENT 3: wire the prompt-prefix record into the GLM engine (audit 253).

`_run` computes the reusable offset, passes it to the runner, and records what the shared State now holds.
Everything is inert while the record is empty (the first request of any conversation), so the cold path and
`sha` are unchanged.

usage: patch_engine_cache.py <engine.py> [--apply]
"""
import ast, sys

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

# 1. import + create the record where the runner is created.
old1 = """            self.runner = Runner(fw, self.limit + self.k + 1, self.k, graphs=GRAPHS, slots=self.parallel)"""
new1 = """            self.runner = Runner(fw, self.limit + self.k + 1, self.k, graphs=GRAPHS, slots=self.parallel)
            self.prefix = PromptPrefix()                 # audit 253: what the shared State holds"""

old2 = """from tensorfold.families.glm_moe_dsa.cuda.runner import Runner"""
if src.count(old2) != 1:
    # find the actual import line for Runner
    import re
    m = re.search(r"^from .*runner import .*Runner.*$", src, re.M)
    if not m:
        print("REFUSED: cannot find the Runner import"); sys.exit(1)
    old2 = m.group(0)
new2 = old2 + "\nfrom tensorfold.families.glm_moe_dsa.cuda.prefix_cache import PromptPrefix"

# 2. _run: decide the offset, pass it, and record the state's new contents.
old3 = """        if self.runner is not None:
            sample = None if s is None else (lambda lg, pos: self._sample(lg, pos, s))
            st = self.runner.generate(prompt, max_tokens, sample, lambda t: stop_eos and t in self.eos, on_tokens, k,
                                      mode, sampling=s)"""
new3 = """        if self.runner is not None:
            sample = None if s is None else (lambda lg, pos: self._sample(lg, pos, s))
            # audit 253: reuse the KV of a prompt prefix the shared State already holds (0 when it holds nothing
            # that this prompt strictly extends). The engine then prefills only the suffix.
            offset = self.prefix.offset_for(prompt)
            st = self.runner.generate(prompt, max_tokens, sample, lambda t: stop_eos and t in self.eos, on_tokens, k,
                                      mode, sampling=s, offset=offset)
            st["cached"] = offset"""

# 3. after the reply, the State holds prompt + its own committed tokens (they are decoded into the cache).
old4 = """            out = st.pop("out")"""
new4 = """            out = st.pop("out")
            self.prefix.record([*prompt, *out])          # audit 253: the State now holds prompt + reply"""

for label, old in (("engine: create the record", old1), ("engine: import", old2),
                   ("engine: offset + pass", old3), ("engine: record", old4)):
    if src.count(old) != 1:
        print("REFUSED: %s: found %d occurrences (need exactly 1)" % (label, src.count(old)))
        sys.exit(1)

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