# -*- coding: utf-8 -*-

"""Incremental decode caches for bit-exact autoregressive generation.

The DoL-OCR RDT has exactly one cross-position operation: causal MLA
attention. Every other sub-layer is per-position (mHC stream mixing, SwiGLU
FFN, RMSNorm) or a sequential recurrence (the NaiveSSM Mamba scan). Therefore
incremental decoding -- appending one token at a time while reusing frozen
earlier state -- is *numerically identical* to re-running the full forward on
the growing prefix, and these caches make that explicit.

Two primitive caches:

* :class:`MLACache` -- post-RoPE ``K``/``V`` of one attention call, grown by
  ``append``. Because each refinement pass of the Stage-2 loop is itself a
  fresh causal attention over the frozen earlier positions, every
  ``(step, layer)`` pair owns its own :class:`MLACache`.
* :class:`MambaCache` -- either the constant-size NaiveSSM recurrent state or
  one call site's isolated upstream ``InferenceParams`` object.

:class:`DecodeCache` is a flat keyed container the orchestrators
(``RDTForCausalLM`` / ``TwoStageCore``) populate lazily; keys encode the call
site (e.g. ``"prelude.0"``, ``"stage1.2"``, ``"stage2.s3.l0"``, ``"coda.1"``).
"""

from __future__ import annotations

from dataclasses import dataclass, field

import torch


@dataclass
class MLACache:
    """Per-attention-call key/value cache (post-RoPE, ``[B, H, T, head_dim]``)."""

    k: torch.Tensor | None = None
    v: torch.Tensor | None = None

    @property
    def length(self) -> int:
        return 0 if self.k is None else self.k.shape[-2]

    def append(
        self, k: torch.Tensor, v: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Append the new tokens' ``k``/``v`` and return the full cached tensors."""

        if self.k is None:
            self.k, self.v = k, v
        else:
            self.k = torch.cat([self.k, k], dim=-2)
            self.v = torch.cat([self.v, v], dim=-2)
        return self.k, self.v


@dataclass
class MambaCache:
    """Per-call-site recurrent state for exactly one Mamba backend.

    ``conv_window`` holds the last ``d_conv - 1`` *pre-convolution* columns
    (shape ``[B, d_inner, d_conv - 1]``); ``ssm_state`` is the selective-scan
    state ``[B, nheads, headdim, d_state]`` in float32. The ``official_*``
    fields hold an isolated upstream cache; they must never be shared between
    logical call sites because upstream indexes state by ``layer_idx``.
    """

    conv_window: torch.Tensor | None = None
    ssm_state: torch.Tensor | None = None
    backend: str | None = None
    official_inference_params: object | None = None
    official_owner: object | None = field(default=None, repr=False)
    official_batch_size: int | None = None
    official_max_seqlen: int | None = None
    official_device: torch.device | None = None
    official_dtype: torch.dtype | None = None
    official_poisoned: bool = False


@dataclass
class DecodeCache:
    """Flat keyed container of per-layer caches for one decode stream."""

    mla: dict[str, MLACache] = field(default_factory=dict)
    mamba: dict[str, MambaCache] = field(default_factory=dict)
    seq_len: int = 0

    def mla_cache(self, key: str) -> MLACache:
        cache = self.mla.get(key)
        if cache is None:
            cache = self.mla[key] = MLACache()
        return cache

    def mamba_cache(self, key: str) -> MambaCache:
        cache = self.mamba.get(key)
        if cache is None:
            cache = self.mamba[key] = MambaCache()
        return cache

    @property
    def past_len(self) -> int:
        """Number of tokens already cached (0 before the prefill call)."""

        return self.seq_len