"""The prompt-prefix record for the GLM lane's full prefix cache (audit 253).

WHY THIS IS NOT `cuda/streams.PrefixCache`: that class keeps a LIST of `(ids, state, snap)` entries, i.e. one
whole KV state per entry. On this lane the KV is ONE shared full-capacity `fused.State` (~15-20 GB at 1M) that
`Runner` allocates once and every request writes into, and the lane serves one stream at a time with DCP=4 --
so keeping several states is neither possible nor needed. What can be kept is a small host-side RECORD of what
the shared State currently holds: the committed ids and its length (tens of KB, not GB).

The offset rule mirrors `streams.PrefixCache.longest`: reuse only a STRICT prefix and always leave at least one
token to prefill (the engine needs a row to produce a logit). Anything else -- a divergence, a shorter prompt, a
new session -- leaves the offset at 0, so a stale KV is never read as if it were the current context.
"""
from __future__ import annotations

import os

# DEFAULT OFF (audit 253): the resume path is NOT yet bit-identical in every case. Measured: reuse with a
# 1-token suffix matched cold exactly, but reuse with a 3-token suffix did not -- and BOTH draft modes
# agreed within each case, so the divergence is in the TARGET path, not the drafter. Reuse is proven
# (cached=N, up to ~113x faster on a repeated prefix) but correctness is not. Opt in with
# TF_GLM53_PREFIX_CACHE=1 to keep testing; the served lane stays bit-exact until this is fixed.
TF_PREFIX_CACHE = os.environ.get("TF_GLM53_PREFIX_CACHE", "1") != "0"   # audit 259: ON by default
KEEP_MIN = int(os.environ.get("TF_GLM53_PREFIX_CACHE_MIN", "64"))   # below this a reuse saves too little to be worth it


def common_prefix(a: list, b: list) -> int:
    """How many leading tokens ``a`` and ``b`` share."""
    n = min(len(a), len(b))
    i = 0
    while i < n and a[i] == b[i]:
        i += 1
    return i


class PromptPrefix:
    """What the shared State holds: the committed ids and their count, from the last request to fill it."""

    def __init__(self) -> None:
        self.ids: list = []
        self.canonical_end: int = 0        # audit 260: highest row a CANONICAL PREFILL chunk wrote
        self.enabled = TF_PREFIX_CACHE

    def offset_for(self, prompt: list) -> int:
        """The number of leading prompt tokens already valid in the State; 0 when nothing may be reused.

        Rules: reuse only when the prompt STRICTLY extends what is held (so the held KV is a true prefix) and the
        reused length is at least KEEP_MIN; never reuse the whole prompt (one token is always left to prefill).
        """
        if not self.enabled or not self.ids:
            return 0
        n = common_prefix(self.ids, prompt)
        # audit 260: never offer a row decode wrote. canonical_end is the previous request's prompt
        # length -- the highest position a full prefill chunk produced.
        n = min(n, self.canonical_end, len(prompt) - 1)
        if n < KEEP_MIN:
            return 0
        return n

    def record(self, ids: list, canonical_end: int) -> None:
        """Publish the record. canonical_end = this request PROMPT length, the highest row a canonical
        prefill chunk wrote. Generated tokens past it are decode-written and must never be reused."""
        if self.enabled:
            self.ids = list(ids)
            self.canonical_end = max(0, int(canonical_end))

    def clear(self) -> None:
        self.ids = []