"""Windowed KV cache for hybrid-attention models (gemma4) under DSpark speculative decoding.

Sliding-attention layers store only `sliding_window + pad` tokens (so 256k context fits in ~5 GB
instead of ~90 GB) while full-attention layers store everything. The `pad` keeps a small speculative
`crop` (rejecting proposal tokens) from eating into the real window. `get_mask_sizes` reports the
true stored length so the model's sliding mask (which still windows to config.sliding_window) aligns.
Validated: 100% argmax match vs the full forward past 1024 tokens, including crop cycles.
"""
import transformers.cache_utils as cu
from transformers.cache_utils import DynamicCache, DynamicLayer


class SpecSlidingLayer(cu.DynamicSlidingWindowLayer):
    def __init__(self, real_window, pad):
        super().__init__(sliding_window=real_window + pad)

    def get_mask_sizes(self, query_length):
        stored = self.keys.shape[-2] if (self.is_initialized and self.keys is not None) else 0
        return stored + query_length, max(self.cumulative_length - stored, 0)

    def crop(self, max_length):
        if max_length < 0:
            max_length = self.cumulative_length + max_length
        remove = self.cumulative_length - max_length
        if remove <= 0:
            return
        n = self.keys.shape[-2]
        self.keys = self.keys[:, :, : n - remove, :]
        self.values = self.values[:, :, : n - remove, :]
        self.cumulative_length = max_length


def build_target_cache(model, pad=64):
    try:
        tc = model.config.get_text_config()
        layer_types = getattr(tc, "layer_types", None)
        sw = getattr(tc, "sliding_window", None)
    except Exception:
        return DynamicCache()
    if not layer_types or not sw:
        return DynamicCache()
    c = DynamicCache()
    c.layers = [SpecSlidingLayer(sw, pad) if lt in ("sliding_attention", "chunked_attention")
                else DynamicLayer() for lt in layer_types]
    return c


def make_draft_cache(draft_model, window, pad=64):
    """Sliding cache for the DSpark draft's accumulated context keys. The draft is full-attention and
    otherwise grows its context cache to the full sequence length; windowing it to the last `window`
    tokens bounds memory. SpecSlidingLayer keeps cumulative_length == absolute position, so the draft's
    position bookkeeping (position_ids[get_seq_length():...]) stays correct. Correctness-safe: the draft
    only proposes; the target verifies every token.
    """
    from transformers.cache_utils import DynamicCache
    try:
        n = draft_model.config.get_text_config().num_hidden_layers
    except Exception:
        n = draft_model.config.num_hidden_layers
    c = DynamicCache()
    c.layers = [SpecSlidingLayer(window, pad) for _ in range(n)]
    return c