"""Cross-attention KV cache for ACE-Step DiT.

Caches encoder K/V projections per layer so they're computed once
(step 0) and reused for steps 1-7. This is a ~30% speedup since
encoder_hidden_states are constant across all denoising steps.

The cache is mutable — pass it by reference and it updates in-place.
No return type changes needed in the call stack.
"""

from typing import Tuple

import mlx.core as mx


class CrossAttentionCache:
    """Stores cross-attention K/V per layer for reuse across denoising steps.

    Usage::

        cache = CrossAttentionCache()
        # Step 0: K/V computed and stored automatically
        # Steps 1-7: K/V retrieved from cache, computation skipped
    """

    def __init__(self):
        self._keys: dict[int, mx.array] = {}
        self._values: dict[int, mx.array] = {}

    def has_layer(self, layer_idx: int) -> bool:
        """Check if K/V for this layer are cached."""
        return layer_idx in self._keys

    def get(self, layer_idx: int) -> Tuple[mx.array, mx.array]:
        """Retrieve cached K/V for a layer.

        Args:
            layer_idx: Layer index.

        Returns:
            (key_states, value_states) both [B, num_kv_heads, S_enc, head_dim]

        Raises:
            KeyError: If layer not yet cached.
        """
        return self._keys[layer_idx], self._values[layer_idx]

    def update(
        self, layer_idx: int, key_states: mx.array, value_states: mx.array
    ) -> None:
        """Store K/V for a layer.

        Args:
            layer_idx: Layer index.
            key_states: [B, num_kv_heads, S_enc, head_dim]
            value_states: [B, num_kv_heads, S_enc, head_dim]
        """
        self._keys[layer_idx] = key_states
        self._values[layer_idx] = value_states

    def reset(self) -> None:
        """Clear all cached K/V."""
        self._keys.clear()
        self._values.clear()

    def __len__(self) -> int:
        return len(self._keys)
