#!/usr/bin/env python3
"""patch_prefix_cache.py -- INCREMENT 2 of the full prefix cache (audit 253).

(a) Drafter.resume(position): continue a context whose ring already holds the positions below `position`.
(b) glm runner: on a resumed prefill, move the drafter to the offset; report `stats["cached"] = offset` so the
    server surfaces it as usage.cached_tokens.

Both are inert at offset=0. Dry-run unless --apply.
usage: patch_prefix_cache.py <dflash2.py> <runner.py> [--apply]
"""
import ast, sys

DF, RUN = sys.argv[1], sys.argv[2]
apply = "--apply" in sys.argv

# ---------------------------------------------------------------- (a) the drafter
df = open(DF).read()
oldA = """    def reset(self) -> None:
        self.context_end = 0
        self.pos_dev.zero_()"""
newA = """    def reset(self) -> None:
        self.context_end = 0
        self.pos_dev.zero_()

    def resume(self, position: int) -> None:
        \"\"\"Continue a context that already holds the positions below ``position`` (audit 253, prefix-cache resume).

        The ring stores each row at ``pos % RING``, and the earlier request of a resumed conversation wrote the
        rows below ``position`` there, so its sliding window is still valid -- only the counters must move. At
        ``offset == 0`` this is a no-op, so the cold path is unchanged.
        \"\"\"
        position = int(position)
        if position <= 0:
            return
        if position > self.capacity:
            raise ValueError(f"resume at {position} is past the drafter's capacity {self.capacity}")
        self.context_end = position
        self.pos_dev.fill_(position)"""

# ---------------------------------------------------------------- (b) the runner
rn = open(RUN).read()
# b1: the drafter must move to the offset before the resumed chunks add their taps.
oldB1 = """        logits = None
        if self.drafter is not None:
            self.drafter.reset()
        if self.cap is not None:
            self.cap.reset()
        for a0, e0 in self.chunks(L0 - offset):"""
newB1 = """        logits = None
        if self.drafter is not None:
            self.drafter.reset()
            if offset:
                self.drafter.resume(offset)      # audit 253: the ring still holds the window below the offset
        if self.cap is not None:
            self.cap.reset()
        for a0, e0 in self.chunks(L0 - offset):"""
# b2: report the reused prefix so the server can surface usage.cached_tokens.
oldB2 = """        if self.dflash:
            return self._generate_dflash(prompt, out, tok, P, max_tokens, sample, stop, on_tokens, prefill_s,
                                         sampling)"""
newB2 = """        if self.dflash:
            st = self._generate_dflash(prompt, out, tok, P, max_tokens, sample, stop, on_tokens, prefill_s,
                                       sampling)
            st["cached"] = offset               # audit 253: prompt tokens served from the kept prefix
            return st"""
oldB3 = """        if self.auto and k:
            return self._generate_auto(out, tok, P, max_tokens, sample, stop, on_tokens, prefill_s, sampling, k)"""
newB3 = """        if self.auto and k:
            st = self._generate_auto(out, tok, P, max_tokens, sample, stop, on_tokens, prefill_s, sampling, k)
            st["cached"] = offset               # audit 253
            return st"""

for label, t, old in (("dflash2: resume()", df, oldA),
                      ("runner: drafter at offset", rn, oldB1),
                      ("runner: cached (dflash)", rn, oldB2),
                      ("runner: cached (auto)", rn, oldB3)):
    if t.count(old) != 1:
        print("REFUSED: %s: found %d occurrences (need exactly 1)" % (label, t.count(old)))
        sys.exit(1)

df2 = df.replace(oldA, newA)
rn2 = rn.replace(oldB1, newB1).replace(oldB2, newB2).replace(oldB3, newB3)
ast.parse(df2); ast.parse(rn2)
print("all 4 anchors unique; both files parse")
if not apply:
    print("dry run -- pass --apply to write"); sys.exit(0)
open(DF, "w").write(df2); open(RUN, "w").write(rn2)
print("APPLIED")