//! Key-Value cache for efficient autoregressive inference.
//!
//! KV caching stores previously computed key and value tensors during generation,
//! avoiding redundant computation of attention for past tokens.
//!
//! ## Memory-Compute Tradeoff
//!
//! - **Without KV cache**: O(n²) attention computation per token
//! - **With KV cache**: O(n) attention computation per token
//!
//! For a sequence of length n, KV caching reduces total generation complexity
//! from O(n³) to O(n²), providing significant speedups for long sequences.
//!
//! ## Tensor Format (SOTA Performance)
//!
//! Keys/values are stored in **attention format** `[B, heads, seq, head_dim]`:
//! - Axis 0: Batch dimension
//! - Axis 1: Number of KV heads
//! - Axis 2: Sequence length (grows during generation)
//! - Axis 3: Head dimension
//!
//! This matches the mlx_lm implementation and eliminates transpose overhead
//! during cached generation. The sequence dimension is axis 2.
//!
//! ## Supported Modes
//!
//! - **Standard KV cache**: Stores all past keys/values (best for short sequences)
//! - **Sliding window cache**: Fixed-size window (constant memory, for long sequences)
//!
//! ## Usage
//!
//! ```ignore
//! let mut cache = KVCache::new(num_layers, max_len, num_kv_heads, head_dim);
//! for step in 0..generation_length {
//!     // Pass keys/values in [B, heads, seq, head_dim] format
//!     let (keys, values) = cache.update_and_fetch(layer_idx, new_keys, new_values)?;
//!     // Use keys, values directly in attention (no transpose needed)
//! }
//! ```

use mlx_rs::{
    Array, Dtype,
    error::Exception,
    ops,
    ops::concatenate_axis,
    ops::indexing::{IndexOp, TryIndexMutOp},
    ops::{dequantize, quantize},
};

/// Configuration for KV cache.
#[derive(Debug, Clone)]
pub struct KVCacheConfig {
    /// Number of transformer layers.
    pub num_layers: usize,
    /// Maximum sequence length.
    pub max_seq_len: usize,
    /// Number of key-value heads (for GQA/MQA).
    pub num_kv_heads: usize,
    /// Dimension per head.
    pub head_dim: usize,
    /// Data type for cached tensors.
    pub dtype: Dtype,
    /// Cache mode.
    pub mode: CacheMode,
    /// Whether to eagerly pre-allocate the full context window upfront.
    /// When true, allocates memory for max_seq_len tokens at creation time.
    /// This provides predictable memory usage but uses more memory initially.
    /// Default: false (lazy allocation in 256-token chunks).
    pub eager_allocate: bool,
    /// Batch size for eager allocation (only used when eager_allocate=true).
    /// Default: 1
    pub eager_batch_size: usize,
}

/// KV cache mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CacheMode {
    /// Standard cache - stores all past tokens.
    Standard,
    /// Sliding window cache with fixed size.
    SlidingWindow {
        /// Maximum number of past tokens to keep in cache.
        window_size: usize,
    },
    /// Rotating cache - circular buffer with fixed max size (MLX-LM parity).
    /// More memory-efficient than sliding window for long sequences.
    Rotating {
        /// Maximum number of tokens to keep.
        max_size: usize,
        /// Number of initial tokens to always keep (typically prompt tokens).
        keep: usize,
    },
    /// Quantized cache - stores K/V in lower precision (MLX-LM parity).
    /// Reduces memory by 2-8x depending on bits.
    Quantized {
        /// Number of bits for quantization (2, 4, or 8).
        bits: u8,
        /// Group size for quantization (default: 64).
        group_size: usize,
    },
}

impl Default for CacheMode {
    fn default() -> Self {
        Self::Standard
    }
}

impl KVCacheConfig {
    /// Create a new KV cache configuration.
    pub fn new(
        num_layers: usize,
        max_seq_len: usize,
        num_kv_heads: usize,
        head_dim: usize,
    ) -> Self {
        Self {
            num_layers,
            max_seq_len,
            num_kv_heads,
            head_dim,
            dtype: Dtype::Float32,
            mode: CacheMode::Standard,
            eager_allocate: false,
            eager_batch_size: 1,
        }
    }

    /// Set the data type for cached tensors.
    pub fn with_dtype(mut self, dtype: Dtype) -> Self {
        self.dtype = dtype;
        self
    }

    /// Set the cache mode.
    pub fn with_mode(mut self, mode: CacheMode) -> Self {
        self.mode = mode;
        self
    }

    /// Enable sliding window mode.
    pub fn with_sliding_window(mut self, window_size: usize) -> Self {
        self.mode = CacheMode::SlidingWindow { window_size };
        self
    }

    /// Enable rotating cache mode (MLX-LM style).
    ///
    /// The rotating cache is a circular buffer that overwrites oldest entries
    /// when full, while optionally preserving `keep` initial tokens.
    ///
    /// # Arguments
    /// * `max_size` - Maximum number of tokens to store
    /// * `keep` - Number of initial tokens to always preserve (0 for none)
    pub fn with_rotating(mut self, max_size: usize, keep: usize) -> Self {
        self.mode = CacheMode::Rotating { max_size, keep };
        self
    }

    /// Enable quantized cache mode (MLX-LM style).
    ///
    /// Stores keys/values in lower precision to reduce memory usage.
    /// - 8-bit: ~2x memory reduction
    /// - 4-bit: ~4x memory reduction
    /// - 2-bit: ~8x memory reduction
    ///
    /// # Arguments
    /// * `bits` - Number of bits (2, 4, or 8)
    /// * `group_size` - Group size for quantization (default: 64)
    pub fn with_quantized(mut self, bits: u8, group_size: usize) -> Self {
        self.mode = CacheMode::Quantized { bits, group_size };
        self
    }

    /// Enable eager pre-allocation of the full context window.
    ///
    /// When enabled, the KV cache will allocate memory for the full `max_seq_len`
    /// at creation time rather than growing dynamically. This provides:
    /// - **Predictable memory usage**: Know exactly how much memory is needed upfront
    /// - **No allocation during generation**: Faster token generation
    /// - **Memory fragmentation prevention**: Single contiguous allocation
    ///
    /// Trade-off: Uses more memory initially even for short sequences.
    ///
    /// # Arguments
    /// * `batch_size` - Batch size to pre-allocate for (typically 1 for inference)
    ///
    /// # Example
    /// ```ignore
    /// let config = KVCacheConfig::new(32, 4096, 8, 128)
    ///     .with_eager_allocate(1);  // Pre-allocate for batch_size=1
    /// let cache = KVCache::new(config);  // ~1GB allocated immediately
    /// ```
    pub fn with_eager_allocate(mut self, batch_size: usize) -> Self {
        self.eager_allocate = true;
        self.eager_batch_size = batch_size;
        self
    }

    /// Calculate the memory footprint for this configuration in bytes.
    ///
    /// Useful for understanding memory requirements before allocation.
    pub fn memory_footprint(&self) -> usize {
        let bytes_per_element = match self.dtype {
            Dtype::Float32 => 4,
            Dtype::Float16 | Dtype::Bfloat16 => 2,
            _ => 4, // Default assumption
        };

        // Per layer: 2 tensors (K, V) × batch × heads × seq × head_dim × bytes
        let per_layer = 2
            * self.eager_batch_size
            * self.num_kv_heads
            * self.max_seq_len
            * self.head_dim
            * bytes_per_element;

        per_layer * self.num_layers
    }

    /// Format the memory footprint as a human-readable string.
    pub fn memory_footprint_human(&self) -> String {
        let bytes = self.memory_footprint();
        if bytes >= 1024 * 1024 * 1024 {
            format!("{:.2} GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0))
        } else if bytes >= 1024 * 1024 {
            format!("{:.2} MB", bytes as f64 / (1024.0 * 1024.0))
        } else if bytes >= 1024 {
            format!("{:.2} KB", bytes as f64 / 1024.0)
        } else {
            format!("{} bytes", bytes)
        }
    }
}

/// Pre-allocation step size in tokens (matches Python mlx-lm).
/// Cache grows in chunks of this size to avoid per-token allocations.
const CACHE_STEP_SIZE: usize = 256;

/// Per-layer KV cache entry.
#[derive(Debug, Clone)]
struct LayerCache {
    /// Cached keys [batch, heads, allocated_seq, head_dim] - attention format for SOTA performance.
    /// Note: allocated_seq may be larger than offset (pre-allocated space).
    keys: Option<Array>,
    /// Cached values [batch, heads, allocated_seq, head_dim] - attention format for SOTA performance.
    values: Option<Array>,
    /// Current offset (actual data length) within the pre-allocated buffer.
    /// This is the number of tokens actually stored, not the buffer size.
    offset: usize,
}

impl LayerCache {
    fn new() -> Self {
        Self {
            keys: None,
            values: None,
            offset: 0,
        }
    }

    /// Create a new layer cache with pre-allocated buffers.
    fn new_eager(
        batch_size: usize,
        num_kv_heads: usize,
        max_seq_len: usize,
        head_dim: usize,
        dtype: Dtype,
    ) -> Result<Self, Exception> {
        let shape = [
            batch_size as i32,
            num_kv_heads as i32,
            max_seq_len as i32,
            head_dim as i32,
        ];
        let keys = Some(ops::zeros_dtype(&shape, dtype)?);
        let values = Some(ops::zeros_dtype(&shape, dtype)?);

        Ok(Self {
            keys,
            values,
            offset: 0,
        })
    }

    fn reset(&mut self) {
        // For eager allocation, just reset offset but keep buffers
        self.offset = 0;
        // Note: We don't set keys/values to None to preserve the pre-allocated buffers
    }

    fn reset_full(&mut self) {
        self.keys = None;
        self.values = None;
        self.offset = 0;
    }
}

/// Key-Value cache for transformer inference.
///
/// Stores computed key and value tensors across generation steps,
/// enabling efficient autoregressive generation.
#[derive(Debug)]
pub struct KVCache {
    /// Configuration.
    config: KVCacheConfig,
    /// Per-layer cache entries.
    layer_caches: Vec<LayerCache>,
    /// Total number of tokens processed.
    total_tokens: usize,
}

impl KVCache {
    /// Create a new KV cache with the given configuration.
    ///
    /// If `eager_allocate` is enabled in the config, this will pre-allocate
    /// the full context window immediately and may fail if memory is insufficient.
    pub fn new(config: KVCacheConfig) -> Self {
        let layer_caches = (0..config.num_layers).map(|_| LayerCache::new()).collect();

        Self {
            config,
            layer_caches,
            total_tokens: 0,
        }
    }

    /// Create a new KV cache with eager pre-allocation.
    ///
    /// Pre-allocates the full `max_seq_len` context window for all layers.
    /// Returns an error if memory allocation fails.
    ///
    /// # Memory Calculation
    /// Memory = 2 × num_layers × batch_size × num_kv_heads × max_seq_len × head_dim × dtype_size
    ///
    /// # Example
    /// ```ignore
    /// let config = KVCacheConfig::new(32, 4096, 8, 128)
    ///     .with_eager_allocate(1)
    ///     .with_dtype(Dtype::Float16);
    /// println!("Will allocate: {}", config.memory_footprint_human());
    /// let cache = KVCache::new_eager(config)?;
    /// ```
    pub fn new_eager(config: KVCacheConfig) -> Result<Self, Exception> {
        if !config.eager_allocate {
            // Fall back to lazy allocation
            return Ok(Self::new(config));
        }

        let mut layer_caches = Vec::with_capacity(config.num_layers);
        for _ in 0..config.num_layers {
            layer_caches.push(LayerCache::new_eager(
                config.eager_batch_size,
                config.num_kv_heads,
                config.max_seq_len,
                config.head_dim,
                config.dtype,
            )?);
        }

        // Evaluate all allocations to materialize them on device
        for cache in &layer_caches {
            if let Some(ref k) = cache.keys {
                k.eval()?;
            }
            if let Some(ref v) = cache.values {
                v.eval()?;
            }
        }

        Ok(Self {
            config,
            layer_caches,
            total_tokens: 0,
        })
    }

    /// Get the cache configuration.
    pub fn config(&self) -> &KVCacheConfig {
        &self.config
    }

    /// Get the current cached sequence length.
    pub fn seq_len(&self) -> usize {
        self.layer_caches.first().map(|c| c.offset).unwrap_or(0)
    }

    /// Get total tokens processed (may differ from seq_len with sliding window).
    pub fn total_tokens(&self) -> usize {
        self.total_tokens
    }

    /// Check if the cache is empty (no tokens stored).
    pub fn is_empty(&self) -> bool {
        // For eager allocation, check offset instead of Option
        self.layer_caches.iter().all(|c| c.offset == 0)
    }

    /// Check if the cache is pre-allocated (eager mode).
    pub fn is_preallocated(&self) -> bool {
        self.config.eager_allocate
            && self
                .layer_caches
                .first()
                .map(|c| c.keys.is_some())
                .unwrap_or(false)
    }

    /// Reset the cache for a new generation.
    ///
    /// For eager-allocated caches, this resets the offset but preserves the
    /// pre-allocated buffers. For lazy caches, this deallocates all memory.
    pub fn reset(&mut self) {
        if self.config.eager_allocate {
            // Eager mode: preserve buffers, just reset offset
            for cache in &mut self.layer_caches {
                cache.reset();
            }
        } else {
            // Lazy mode: deallocate to free memory
            for cache in &mut self.layer_caches {
                cache.reset_full();
            }
        }
        self.total_tokens = 0;
    }

    /// Fully reset the cache, deallocating all memory.
    ///
    /// Unlike `reset()`, this always deallocates the buffers even for
    /// eager-allocated caches. Use this to free memory when done.
    pub fn reset_full(&mut self) {
        for cache in &mut self.layer_caches {
            cache.reset_full();
        }
        self.total_tokens = 0;
    }

    /// Update the cache with new keys and values for a layer.
    ///
    /// Keys/values are expected in **attention format** `[B, heads, seq, head_dim]`
    /// where axis 2 is the sequence dimension. This matches mlx_lm for SOTA performance.
    ///
    /// Uses pre-allocation in chunks of CACHE_STEP_SIZE (256 tokens) and in-place
    /// slice assignment for O(1) amortized per-token updates (matching Python mlx-lm).
    ///
    /// # Arguments
    /// * `layer_idx` - Layer index
    /// * `new_keys` - New keys [batch, heads, new_seq, head_dim]
    /// * `new_values` - New values [batch, heads, new_seq, head_dim]
    ///
    /// # Returns
    /// (cached_keys, cached_values) - Slice views of cached keys/values
    pub fn update_and_fetch(
        &mut self,
        layer_idx: usize,
        new_keys: &Array,
        new_values: &Array,
    ) -> Result<(Array, Array), Exception> {
        if layer_idx >= self.config.num_layers {
            return Err(Exception::custom(format!(
                "Layer index {} out of range (num_layers={})",
                layer_idx, self.config.num_layers
            )));
        }

        let cache = &mut self.layer_caches[layer_idx];
        // Sequence dimension is axis 2 in [B, heads, seq, head_dim] format
        let new_seq_len = new_keys.dim(2) as usize;
        let prev_offset = cache.offset;

        // Update total tokens count (only count for first layer to avoid double counting)
        if layer_idx == 0 {
            self.total_tokens += new_seq_len;
        }

        // Check if we need to grow the cache (pre-allocation pattern from mlx_lm)
        let needs_growth = cache.keys.is_none() || {
            let allocated = cache.keys.as_ref().unwrap().dim(2) as usize;
            prev_offset + new_seq_len > allocated
        };

        if needs_growth {
            // Pre-allocate in chunks of CACHE_STEP_SIZE
            let n_steps = (CACHE_STEP_SIZE + new_seq_len - 1) / CACHE_STEP_SIZE;
            let new_alloc_len = n_steps * CACHE_STEP_SIZE;

            // Get shape from new_keys [B, heads, _, head_dim]
            let batch = new_keys.dim(0);
            let heads = new_keys.dim(1);
            let head_dim = new_keys.dim(3);

            // Create new zero-filled buffer
            let k_shape = [batch, heads, new_alloc_len as i32, head_dim];
            let v_shape = [batch, heads, new_alloc_len as i32, head_dim];
            let new_k_buffer = ops::zeros_dtype(&k_shape, new_keys.dtype())?;
            let new_v_buffer = ops::zeros_dtype(&v_shape, new_values.dtype())?;

            if let (Some(existing_k), Some(existing_v)) = (&cache.keys, &cache.values) {
                // Concatenate existing data with new buffer
                cache.keys = Some(concatenate_axis(&[existing_k, &new_k_buffer], 2)?);
                cache.values = Some(concatenate_axis(&[existing_v, &new_v_buffer], 2)?);
            } else {
                cache.keys = Some(new_k_buffer);
                cache.values = Some(new_v_buffer);
            }
        }

        // Update offset before in-place assignment
        cache.offset = prev_offset + new_seq_len;

        // In-place slice assignment: cache[..., prev:offset, :] = new_keys
        // Use TryIndexMutOp for O(1) update instead of concatenate
        let k_buf = cache.keys.as_mut().unwrap();
        let v_buf = cache.values.as_mut().unwrap();

        // Assign new keys/values into the pre-allocated buffer
        k_buf.try_index_mut(
            (.., .., prev_offset as i32..cache.offset as i32, ..),
            new_keys,
        )?;
        v_buf.try_index_mut(
            (.., .., prev_offset as i32..cache.offset as i32, ..),
            new_values,
        )?;

        // Apply cache mode limits (sliding window, rotating, etc.)
        let final_offset = match self.config.mode {
            CacheMode::SlidingWindow { window_size } => {
                if cache.offset > window_size {
                    // For sliding window, we need to shift data and adjust offset
                    // This is a less common path, so we can do a copy here
                    let shift = cache.offset - window_size;
                    let k = cache.keys.as_ref().unwrap();
                    let v = cache.values.as_ref().unwrap();
                    cache.keys = Some(k.index((.., .., shift as i32..cache.offset as i32, ..)));
                    cache.values = Some(v.index((.., .., shift as i32..cache.offset as i32, ..)));
                    cache.offset = window_size;
                }
                cache.offset
            }
            CacheMode::Rotating { max_size, .. } => {
                if cache.offset > max_size {
                    let shift = cache.offset - max_size;
                    let k = cache.keys.as_ref().unwrap();
                    let v = cache.values.as_ref().unwrap();
                    cache.keys = Some(k.index((.., .., shift as i32..cache.offset as i32, ..)));
                    cache.values = Some(v.index((.., .., shift as i32..cache.offset as i32, ..)));
                    cache.offset = max_size;
                }
                cache.offset
            }
            CacheMode::Quantized { .. } | CacheMode::Standard => {
                if cache.offset > self.config.max_seq_len {
                    return Err(Exception::custom(format!(
                        "KV cache exceeded max_seq_len: {} > {}",
                        cache.offset, self.config.max_seq_len
                    )));
                }
                cache.offset
            }
        };

        // Return slice views (not clones) - matches Python mlx_lm pattern
        let k = cache.keys.as_ref().unwrap();
        let v = cache.values.as_ref().unwrap();
        Ok((
            k.index((.., .., ..final_offset as i32, ..)),
            v.index((.., .., ..final_offset as i32, ..)),
        ))
    }

    /// Legacy update method - DEPRECATED, use update_and_fetch instead.
    /// Kept for backwards compatibility.
    #[deprecated(
        since = "0.2.0",
        note = "Use update_and_fetch instead for SOTA performance"
    )]
    pub fn update(
        &mut self,
        layer_idx: usize,
        new_keys: &Array,
        new_values: &Array,
    ) -> Result<(Array, Array), Exception> {
        self.update_and_fetch(layer_idx, new_keys, new_values)
    }

    /// Get cached keys and values for a layer without updating.
    ///
    /// Returns sliced views up to the actual sequence length (not the pre-allocated buffer size).
    /// Returns None if the cache for this layer is empty.
    pub fn get(&self, layer_idx: usize) -> Option<(Array, Array)> {
        let cache = self.layer_caches.get(layer_idx)?;
        match (&cache.keys, &cache.values) {
            (Some(k), Some(v)) if cache.offset > 0 => {
                // Return sliced view up to actual offset (not full pre-allocated buffer)
                Some((
                    k.index((.., .., ..cache.offset as i32, ..)),
                    v.index((.., .., ..cache.offset as i32, ..)),
                ))
            }
            _ => None,
        }
    }

    /// Get the offset for RoPE when using cache.
    ///
    /// This is the starting position for new tokens when computing
    /// rotary embeddings during cached generation.
    pub fn rope_offset(&self) -> i32 {
        self.seq_len() as i32
    }

    /// Estimate memory usage of the current cache in bytes.
    pub fn memory_usage(&self) -> usize {
        let mut total = 0;
        for cache in &self.layer_caches {
            if let Some(keys) = &cache.keys {
                // Approximate: shape elements * dtype size
                let elements: usize = keys.shape().iter().map(|&d| d as usize).product();
                total += elements * dtype_size(self.config.dtype);
            }
            if let Some(values) = &cache.values {
                let elements: usize = values.shape().iter().map(|&d| d as usize).product();
                total += elements * dtype_size(self.config.dtype);
            }
        }
        total
    }

    /// Estimate maximum memory usage for the configured cache.
    pub fn max_memory_usage(&self) -> usize {
        let seq_len = match self.config.mode {
            CacheMode::SlidingWindow { window_size } => window_size,
            CacheMode::Rotating { max_size, .. } => max_size,
            CacheMode::Quantized { .. } | CacheMode::Standard => self.config.max_seq_len,
        };

        // batch=1 assumed, 2 tensors (K,V) per layer
        let elements_per_layer = seq_len * self.config.num_kv_heads * self.config.head_dim;
        let bytes_per_layer = elements_per_layer * dtype_size(self.config.dtype) * 2;

        self.config.num_layers * bytes_per_layer
    }
}

/// Helper to get dtype size in bytes.
fn dtype_size(dtype: Dtype) -> usize {
    match dtype {
        Dtype::Float32 => 4,
        Dtype::Float64 => 8,
        Dtype::Float16 | Dtype::Bfloat16 => 2,
        Dtype::Int32 => 4,
        Dtype::Int64 => 8,
        Dtype::Int16 => 2,
        Dtype::Int8 | Dtype::Uint8 => 1,
        Dtype::Uint16 => 2,
        Dtype::Uint32 => 4,
        Dtype::Uint64 => 8,
        Dtype::Bool => 1,
        Dtype::Complex64 => 8,
    }
}

/// Convenience function to create a standard KV cache.
pub fn create_kv_cache(
    num_layers: usize,
    max_seq_len: usize,
    num_kv_heads: usize,
    head_dim: usize,
) -> KVCache {
    KVCache::new(KVCacheConfig::new(
        num_layers,
        max_seq_len,
        num_kv_heads,
        head_dim,
    ))
}

/// Convenience function to create a sliding window KV cache.
pub fn create_sliding_window_cache(
    num_layers: usize,
    window_size: usize,
    num_kv_heads: usize,
    head_dim: usize,
) -> KVCache {
    KVCache::new(
        KVCacheConfig::new(num_layers, window_size * 2, num_kv_heads, head_dim)
            .with_sliding_window(window_size),
    )
}

/// Batch of KV caches for parallel generation.
#[derive(Debug)]
pub struct BatchKVCache {
    /// Individual caches per batch entry.
    caches: Vec<KVCache>,
}

impl BatchKVCache {
    /// Create a new batch of KV caches.
    pub fn new(batch_size: usize, config: KVCacheConfig) -> Self {
        let caches = (0..batch_size)
            .map(|_| KVCache::new(config.clone()))
            .collect();
        Self { caches }
    }

    /// Get the batch size.
    pub fn batch_size(&self) -> usize {
        self.caches.len()
    }

    /// Get a mutable reference to a specific cache.
    pub fn get_mut(&mut self, batch_idx: usize) -> Option<&mut KVCache> {
        self.caches.get_mut(batch_idx)
    }

    /// Reset all caches.
    pub fn reset(&mut self) {
        for cache in &mut self.caches {
            cache.reset();
        }
    }

    /// Reset specific cache entries (for finished sequences).
    pub fn reset_indices(&mut self, indices: &[usize]) {
        for &idx in indices {
            if let Some(cache) = self.caches.get_mut(idx) {
                cache.reset();
            }
        }
    }

    /// Total memory usage across all caches.
    pub fn memory_usage(&self) -> usize {
        self.caches.iter().map(|c| c.memory_usage()).sum()
    }
}

// =============================================================================
// RotatingKVCache - MLX-LM Parity Implementation
// =============================================================================

/// Rotating KV cache that acts as a circular buffer.
///
/// This implementation matches MLX-LM's `RotatingKVCache` for full parity.
/// It maintains a fixed-size buffer that overwrites oldest entries when full,
/// while optionally preserving initial "keep" tokens (typically prompt tokens).
///
/// # Memory Efficiency
///
/// Unlike standard KV cache which grows unbounded, rotating cache maintains
/// constant memory usage regardless of generation length:
/// - Memory = max_size * num_heads * head_dim * 2 (K+V) * dtype_size
///
/// # Temporal Ordering
///
/// The cache internally tracks a write index (`_idx`) that wraps around.
/// When reading, entries are reordered to maintain temporal consistency
/// for attention computation.
#[derive(Debug)]
pub struct RotatingKVCache {
    /// Cached keys [batch, heads, seq, head_dim].
    keys: Option<Array>,
    /// Cached values [batch, heads, seq, head_dim].
    values: Option<Array>,
    /// Total offset (tokens seen, may exceed max_size).
    offset: usize,
    /// Current write index in circular buffer.
    _idx: usize,
    /// Maximum size of the cache.
    max_size: usize,
    /// Number of initial tokens to always preserve.
    keep: usize,
    /// Allocation step size.
    step: usize,
}

impl RotatingKVCache {
    /// Create a new rotating KV cache.
    ///
    /// # Arguments
    /// * `max_size` - Maximum number of tokens to store
    /// * `keep` - Number of initial tokens to always preserve (default: 0)
    pub fn new(max_size: usize, keep: usize) -> Self {
        Self {
            keys: None,
            values: None,
            offset: 0,
            _idx: 0,
            max_size,
            keep,
            step: 256,
        }
    }

    /// Get the current length of the cache (capped at max_size).
    pub fn len(&self) -> usize {
        self.offset.min(self.max_size)
    }

    /// Check if the cache is empty.
    pub fn is_empty(&self) -> bool {
        self.offset == 0
    }

    /// Get total tokens processed.
    pub fn offset(&self) -> usize {
        self.offset
    }

    /// Get the RoPE offset for position encoding.
    pub fn rope_offset(&self) -> i32 {
        self.offset as i32
    }

    /// Reset the cache.
    pub fn reset(&mut self) {
        self.keys = None;
        self.values = None;
        self.offset = 0;
        self._idx = 0;
    }

    /// Trim entries from the front of the cache.
    fn trim_impl(
        &self,
        trim_size: usize,
        v: &Array,
        append: Option<&Array>,
    ) -> Result<Array, Exception> {
        if trim_size > 0 {
            // Keep initial tokens, then slice from after trim
            let kept = v.index((.., .., ..self.keep as i32, ..));
            let rest = v.index((.., .., (trim_size + self.keep) as i32.., ..));

            let parts: Vec<&Array> = if let Some(a) = append {
                vec![&kept, &rest, a]
            } else {
                vec![&kept, &rest]
            };
            concatenate_axis(&parts, 2)
        } else if let Some(a) = append {
            concatenate_axis(&[v, a], 2)
        } else {
            Ok(v.clone())
        }
    }

    /// Reorder cache into temporal order (for reading).
    fn temporal_order(&self, v: &Array) -> Array {
        let cache_len = v.dim(2) as usize;

        if self._idx == cache_len {
            // No wrap-around yet
            v.clone()
        } else if self._idx < self.offset {
            // Wrapped around: reorder [keep][idx..][keep..idx]
            let kept = v.index((.., .., ..self.keep as i32, ..));
            let after_idx = v.index((.., .., self._idx as i32.., ..));
            let before_idx = v.index((.., .., self.keep as i32..self._idx as i32, ..));

            concatenate_axis(&[&kept, &after_idx, &before_idx], 2).unwrap_or_else(|_| v.clone())
        } else {
            // Not full yet
            v.index((.., .., ..self._idx as i32, ..))
        }
    }

    /// Update with concatenation (for multi-token updates).
    fn update_concat(&mut self, keys: &Array, values: &Array) -> Result<(Array, Array), Exception> {
        let num_steps = keys.dim(2) as usize;

        if self.keys.is_none() {
            self.keys = Some(keys.clone());
            self.values = Some(values.clone());
        } else {
            // Put in temporal order
            let ordered_k = self.temporal_order(self.keys.as_ref().unwrap());
            let ordered_v = self.temporal_order(self.values.as_ref().unwrap());
            self._idx = ordered_k.dim(2) as usize;

            // Trim to maintain max_size
            let trim_size = self._idx.saturating_sub(self.max_size - 1);
            self.keys = Some(self.trim_impl(trim_size, &ordered_k, Some(keys))?);
            self.values = Some(self.trim_impl(trim_size, &ordered_v, Some(values))?);
        }

        self.offset += num_steps;
        self._idx = self.keys.as_ref().unwrap().dim(2) as usize;

        Ok((
            self.keys.as_ref().unwrap().clone(),
            self.values.as_ref().unwrap().clone(),
        ))
    }

    /// Update in-place (for single-token updates).
    fn update_in_place(
        &mut self,
        keys: &Array,
        values: &Array,
    ) -> Result<(Array, Array), Exception> {
        let batch = keys.dim(0) as usize;
        let n_kv_heads = keys.dim(1) as usize;
        let num_steps = keys.dim(2) as usize;
        let k_head_dim = keys.dim(3) as usize;
        let v_head_dim = values.dim(3) as usize;

        // Grow cache if needed
        let needs_growth = self.keys.is_none()
            || (self.offset >= self.keys.as_ref().unwrap().dim(2) as usize
                && (self.keys.as_ref().unwrap().dim(2) as usize) < self.max_size);
        if needs_growth {
            let new_size = self.step.min(self.max_size - self.offset);
            let k_shape = [
                batch as i32,
                n_kv_heads as i32,
                new_size as i32,
                k_head_dim as i32,
            ];
            let v_shape = [
                batch as i32,
                n_kv_heads as i32,
                new_size as i32,
                v_head_dim as i32,
            ];

            let new_k = Array::zeros::<f32>(&k_shape)?;
            let new_v = Array::zeros::<f32>(&v_shape)?;

            if let Some(ref existing_k) = self.keys {
                self.keys = Some(concatenate_axis(&[existing_k, &new_k], 2)?);
                self.values = Some(concatenate_axis(
                    &[self.values.as_ref().unwrap(), &new_v],
                    2,
                )?);
            } else {
                self.keys = Some(new_k);
                self.values = Some(new_v);
            }
            self._idx = self.offset;
        }

        // Trim if exceeding max_size
        let cache_len = self.keys.as_ref().unwrap().dim(2) as usize;
        if cache_len > self.max_size {
            let trim_size = cache_len - self.max_size;
            let trimmed_k = self.trim_impl(trim_size, self.keys.as_ref().unwrap(), None)?;
            let trimmed_v = self.trim_impl(trim_size, self.values.as_ref().unwrap(), None)?;
            self.keys = Some(trimmed_k);
            self.values = Some(trimmed_v);
            self._idx = self.max_size;
        }

        // Rotate if at max
        if self._idx == self.max_size {
            self._idx = self.keep;
        }

        // In-place update using scatter or slice assignment
        // For now, we reconstruct the array (MLX doesn't have true in-place slice assignment in Rust bindings)
        let k = self.keys.as_ref().unwrap();
        let v = self.values.as_ref().unwrap();

        // Build updated cache by concatenating before, new, and after
        let before_k = if self._idx > 0 {
            Some(k.index((.., .., ..self._idx as i32, ..)))
        } else {
            None
        };
        let after_k = if self._idx + num_steps < k.dim(2) as usize {
            Some(k.index((.., .., (self._idx + num_steps) as i32.., ..)))
        } else {
            None
        };

        let before_v = if self._idx > 0 {
            Some(v.index((.., .., ..self._idx as i32, ..)))
        } else {
            None
        };
        let after_v = if self._idx + num_steps < v.dim(2) as usize {
            Some(v.index((.., .., (self._idx + num_steps) as i32.., ..)))
        } else {
            None
        };

        // Assemble new cache
        let mut k_parts: Vec<&Array> = Vec::new();
        if let Some(ref bk) = before_k {
            k_parts.push(bk);
        }
        k_parts.push(keys);
        if let Some(ref ak) = after_k {
            k_parts.push(ak);
        }

        let mut v_parts: Vec<&Array> = Vec::new();
        if let Some(ref bv) = before_v {
            v_parts.push(bv);
        }
        v_parts.push(values);
        if let Some(ref av) = after_v {
            v_parts.push(av);
        }

        self.keys = Some(concatenate_axis(&k_parts, 2)?);
        self.values = Some(concatenate_axis(&v_parts, 2)?);

        self.offset += num_steps;
        self._idx += num_steps;

        // Return slice if not full yet
        if self.offset < self.max_size {
            Ok((
                self.keys
                    .as_ref()
                    .unwrap()
                    .index((.., .., ..self.offset as i32, ..)),
                self.values
                    .as_ref()
                    .unwrap()
                    .index((.., .., ..self.offset as i32, ..)),
            ))
        } else {
            Ok((
                self.keys.as_ref().unwrap().clone(),
                self.values.as_ref().unwrap().clone(),
            ))
        }
    }

    /// Update cache with new keys and values.
    ///
    /// Uses in-place update for single tokens, concatenation for multi-token.
    pub fn update_and_fetch(
        &mut self,
        keys: &Array,
        values: &Array,
    ) -> Result<(Array, Array), Exception> {
        if keys.dim(2) == 1 {
            self.update_in_place(keys, values)
        } else {
            self.update_concat(keys, values)
        }
    }

    /// Check if cache can be trimmed.
    pub fn is_trimmable(&self) -> bool {
        self.offset < self.max_size
    }

    /// Trim n tokens from the cache.
    pub fn trim(&mut self, n: usize) -> usize {
        let trimmed = n.min(self.offset);
        self.offset -= trimmed;
        self._idx = self._idx.saturating_sub(trimmed);
        trimmed
    }
}

// =============================================================================
// QuantizedKVCache - MLX-LM Parity Implementation
// =============================================================================

/// Quantized representation of cached K/V tensors.
#[derive(Debug, Clone)]
struct QuantizedTensor {
    /// Quantized data (packed integers).
    data: Array,
    /// Scale factors per group.
    scales: Array,
    /// Bias/zero-points per group.
    biases: Array,
}

/// Quantized KV cache that stores keys/values in lower precision.
///
/// This implementation matches MLX-LM's `QuantizedKVCache` for full parity.
/// It reduces memory usage significantly while maintaining acceptable quality:
/// - 8-bit: ~50% memory reduction
/// - 4-bit: ~75% memory reduction
///
/// # Quantization Scheme
///
/// Uses block-wise quantization with configurable group size:
/// - Each group of `group_size` elements shares a scale and bias
/// - Values are quantized as: `quantized = round((value - bias) / scale)`
/// - Dequantized as: `value = quantized * scale + bias`
///
/// # Note
///
/// Requires MLX to have quantize/dequantize operations available.
/// Falls back to standard cache if quantization fails.
#[derive(Debug)]
pub struct QuantizedKVCache {
    /// Quantized keys.
    keys: Option<QuantizedTensor>,
    /// Quantized values.
    values: Option<QuantizedTensor>,
    /// Total offset (tokens seen).
    offset: usize,
    /// Number of bits for quantization.
    bits: u8,
    /// Group size for quantization.
    group_size: usize,
    /// Allocation step size.
    step: usize,
    /// Original dtype for dequantization.
    dtype: Dtype,
}

impl QuantizedKVCache {
    /// Create a new quantized KV cache.
    ///
    /// # Arguments
    /// * `bits` - Number of bits (2, 4, or 8)
    /// * `group_size` - Group size for quantization (default: 64)
    pub fn new(bits: u8, group_size: usize) -> Self {
        assert!(
            bits == 2 || bits == 4 || bits == 8,
            "bits must be 2, 4, or 8"
        );
        Self {
            keys: None,
            values: None,
            offset: 0,
            bits,
            group_size,
            step: 256,
            dtype: Dtype::Float16, // Default to f16
        }
    }

    /// Get the current length.
    pub fn len(&self) -> usize {
        self.offset
    }

    /// Check if empty.
    pub fn is_empty(&self) -> bool {
        self.offset == 0
    }

    /// Get RoPE offset.
    pub fn rope_offset(&self) -> i32 {
        self.offset as i32
    }

    /// Reset the cache.
    pub fn reset(&mut self) {
        self.keys = None;
        self.values = None;
        self.offset = 0;
    }

    /// Pack a 4-D float tensor `[B, H, S, D]` into a [`QuantizedTensor`].
    ///
    /// MLX's `quantize` operates on 2-D matrices where groups are formed along
    /// the last dimension.  The strategy is:
    ///
    /// 1. Cast to float32 (quantize requires a floating-point input).
    /// 2. Reshape `[B, H, S, D]` → `[B*H*S, D]` so the last axis is the head
    ///    dimension, which is where the group structure lives.
    /// 3. Invoke `mlx_rs::ops::quantize`, which returns
    ///    `(w_q [rows, D/el_per_int], scales [rows, D/group_size], biases [rows, D/group_size])`.
    /// 4. Reshape each of those 2-D results back to 4-D so they can be
    ///    concatenated across the sequence dimension later.
    ///
    /// # Panics / Errors
    ///
    /// Returns an `Exception` if MLX rejects the shapes (e.g., `D` not
    /// divisible by `group_size`).  The caller is responsible for ensuring the
    /// head dimension satisfies this constraint or for padding prior to calling.
    fn quantize_tensor(&self, tensor: &Array) -> Result<QuantizedTensor, Exception> {
        let shape = tensor.shape();
        let batch = shape[0] as usize;
        let heads = shape[1] as usize;
        let seq = shape[2] as usize;
        let dim = shape[3] as usize;

        // MLX quantize requires a float32 input.
        let float_tensor = tensor.as_type::<f32>()?;

        // Collapse the three leading dimensions into one so the last axis is
        // the head dimension that we want to quantize over.
        let rows = (batch * heads * seq) as i32;
        let flat = float_tensor.reshape(&[rows, dim as i32])?;

        // mlx_rs::ops::quantize returns (w_q, scales, biases).
        //   w_q    : [rows, dim * bits / 32]   (packed u32)
        //   scales : [rows, dim / group_size]
        //   biases : [rows, dim / group_size]
        let group_size_i32 = self.group_size as i32;
        let bits_i32 = self.bits as i32;
        let (w_q, scales_2d, biases_2d) = quantize(&flat, group_size_i32, bits_i32)?;

        // Reshape packed data back to [B, H, S, packed_dim].
        let packed_dim = w_q.dim(1);
        let data = w_q.reshape(&[batch as i32, heads as i32, seq as i32, packed_dim])?;

        // Reshape scales/biases back to [B, H, S, num_groups].
        let num_groups = scales_2d.dim(1);
        let scales = scales_2d.reshape(&[batch as i32, heads as i32, seq as i32, num_groups])?;
        let biases = biases_2d.reshape(&[batch as i32, heads as i32, seq as i32, num_groups])?;

        Ok(QuantizedTensor {
            data,
            scales,
            biases,
        })
    }

    /// Unpack a [`QuantizedTensor`] back into a float tensor `[B, H, S, D]`.
    ///
    /// This is the exact inverse of [`Self::quantize_tensor`]:
    ///
    /// 1. Flatten the 4-D packed data and metadata to 2-D.
    /// 2. Invoke `mlx_rs::ops::dequantize`, which reconstructs a float32
    ///    matrix `[rows, D]`.
    /// 3. Reshape back to `[B, H, S, D]` and cast to the original dtype that
    ///    was fed in (stored in `self.dtype`).
    fn dequantize_tensor(&self, qtensor: &QuantizedTensor) -> Result<Array, Exception> {
        let shape = qtensor.data.shape();
        let batch = shape[0] as usize;
        let heads = shape[1] as usize;
        let seq = shape[2] as usize;
        // packed_dim = D * bits / 32 => D = packed_dim * 32 / bits
        let packed_dim = shape[3] as usize;
        let el_per_int = 32usize / self.bits as usize;
        let dim = packed_dim * el_per_int;

        let rows = (batch * heads * seq) as i32;

        // Flatten to 2D for the MLX op.
        let flat_data = qtensor.data.reshape(&[rows, packed_dim as i32])?;
        let flat_scales = qtensor.scales.reshape(&[rows, qtensor.scales.dim(3)])?;
        let flat_biases = qtensor.biases.reshape(&[rows, qtensor.biases.dim(3)])?;

        // dequantize returns a float32 array [rows, dim].
        let group_size_i32 = self.group_size as i32;
        let bits_i32 = self.bits as i32;
        let flat_float = dequantize(
            &flat_data,
            &flat_scales,
            &flat_biases,
            group_size_i32,
            bits_i32,
        )?;

        // Restore 4-D layout [B, H, S, D] and cast to the original dtype.
        let out_4d = flat_float.reshape(&[batch as i32, heads as i32, seq as i32, dim as i32])?;

        // Cast back to the dtype we received (typically f16 / bf16).
        out_4d.as_dtype(self.dtype)
    }

    /// Update cache with new keys and values.
    pub fn update_and_fetch(
        &mut self,
        keys: &Array,
        values: &Array,
    ) -> Result<(Array, Array), Exception> {
        self.dtype = keys.dtype();
        let num_steps = keys.dim(2) as usize;

        // Quantize new keys/values
        let q_keys = self.quantize_tensor(keys)?;
        let q_values = self.quantize_tensor(values)?;

        if self.keys.is_none() {
            self.keys = Some(q_keys);
            self.values = Some(q_values);
        } else {
            // Concatenate quantized tensors
            let existing_k = self.keys.as_ref().unwrap();
            let existing_v = self.values.as_ref().unwrap();

            self.keys = Some(QuantizedTensor {
                data: concatenate_axis(&[&existing_k.data, &q_keys.data], 2)?,
                scales: concatenate_axis(&[&existing_k.scales, &q_keys.scales], 2)?,
                biases: concatenate_axis(&[&existing_k.biases, &q_keys.biases], 2)?,
            });
            self.values = Some(QuantizedTensor {
                data: concatenate_axis(&[&existing_v.data, &q_values.data], 2)?,
                scales: concatenate_axis(&[&existing_v.scales, &q_values.scales], 2)?,
                biases: concatenate_axis(&[&existing_v.biases, &q_values.biases], 2)?,
            });
        }

        self.offset += num_steps;

        // Dequantize for attention computation
        let dk = self.dequantize_tensor(self.keys.as_ref().unwrap())?;
        let dv = self.dequantize_tensor(self.values.as_ref().unwrap())?;

        Ok((dk, dv))
    }

    /// Check if trimmable.
    pub fn is_trimmable(&self) -> bool {
        true
    }

    /// Trim n tokens.
    pub fn trim(&mut self, n: usize) -> usize {
        let trimmed = n.min(self.offset);
        self.offset -= trimmed;
        trimmed
    }

    /// Estimated memory usage.
    pub fn memory_usage(&self) -> usize {
        if let Some(ref k) = self.keys {
            let k_elements: usize = k.data.shape().iter().map(|&d| d as usize).product();
            let s_elements: usize = k.scales.shape().iter().map(|&d| d as usize).product();
            // data uses 4 bytes (u32), scales/biases use 2 bytes each (f16)
            let k_bytes = k_elements * 4 + s_elements * 4;
            k_bytes * 2 // K + V
        } else {
            0
        }
    }
}

/// Convenience function to create a rotating KV cache.
pub fn create_rotating_cache(max_size: usize, keep: usize) -> RotatingKVCache {
    RotatingKVCache::new(max_size, keep)
}

/// Convenience function to create a quantized KV cache.
pub fn create_quantized_cache(bits: u8, group_size: usize) -> QuantizedKVCache {
    QuantizedKVCache::new(bits, group_size)
}

// =============================================================================
// PagedKVCache - Block-based KV cache for efficient batched inference
// =============================================================================

/// Block size for paged attention (tokens per block).
/// 32 tokens is optimal for Apple Silicon (matches GPU cache lines).
pub const DEFAULT_BLOCK_SIZE: usize = 32;

/// Configuration for paged KV cache.
#[derive(Debug, Clone)]
pub struct PagedKVCacheConfig {
    /// Number of transformer layers.
    pub num_layers: usize,
    /// Number of key-value heads.
    pub num_kv_heads: usize,
    /// Head dimension.
    pub head_dim: usize,
    /// Block size (tokens per block).
    pub block_size: usize,
    /// Maximum number of blocks to allocate.
    pub max_blocks: usize,
    /// Data type for cached tensors.
    pub dtype: Dtype,
}

impl PagedKVCacheConfig {
    /// Create a new paged KV cache configuration.
    ///
    /// # Arguments
    /// * `num_layers` - Number of transformer layers
    /// * `num_kv_heads` - Number of KV heads
    /// * `head_dim` - Dimension per head
    /// * `max_seq_len` - Maximum sequence length to support
    pub fn new(
        num_layers: usize,
        num_kv_heads: usize,
        head_dim: usize,
        max_seq_len: usize,
    ) -> Self {
        let max_blocks = (max_seq_len + DEFAULT_BLOCK_SIZE - 1) / DEFAULT_BLOCK_SIZE;
        Self {
            num_layers,
            num_kv_heads,
            head_dim,
            block_size: DEFAULT_BLOCK_SIZE,
            max_blocks,
            dtype: Dtype::Float16,
        }
    }

    /// Set the block size.
    pub fn with_block_size(mut self, block_size: usize) -> Self {
        self.block_size = block_size;
        self.max_blocks = (self.max_blocks * DEFAULT_BLOCK_SIZE + block_size - 1) / block_size;
        self
    }

    /// Set the dtype.
    pub fn with_dtype(mut self, dtype: Dtype) -> Self {
        self.dtype = dtype;
        self
    }

    /// Set the maximum number of blocks.
    pub fn with_max_blocks(mut self, max_blocks: usize) -> Self {
        self.max_blocks = max_blocks;
        self
    }
}

/// Block allocator for managing physical memory blocks.
#[derive(Debug)]
pub struct BlockAllocator {
    /// Free block indices.
    free_blocks: Vec<usize>,
    /// Total number of blocks allocated.
    total_blocks: usize,
    /// Block size in tokens.
    block_size: usize,
}

impl BlockAllocator {
    /// Create a new block allocator.
    pub fn new(num_blocks: usize, block_size: usize) -> Self {
        Self {
            free_blocks: (0..num_blocks).rev().collect(), // Stack-like for LIFO reuse
            total_blocks: num_blocks,
            block_size,
        }
    }

    /// Allocate a block, returning its index.
    pub fn allocate(&mut self) -> Option<usize> {
        self.free_blocks.pop()
    }

    /// Allocate multiple blocks.
    pub fn allocate_n(&mut self, n: usize) -> Option<Vec<usize>> {
        if self.free_blocks.len() < n {
            return None;
        }
        let blocks: Vec<usize> = (0..n).map(|_| self.free_blocks.pop().unwrap()).collect();
        Some(blocks)
    }

    /// Free a block.
    pub fn free(&mut self, block_idx: usize) {
        debug_assert!(block_idx < self.total_blocks);
        self.free_blocks.push(block_idx);
    }

    /// Free multiple blocks.
    pub fn free_all(&mut self, blocks: &[usize]) {
        for &block_idx in blocks {
            self.free(block_idx);
        }
    }

    /// Get the number of free blocks.
    pub fn num_free(&self) -> usize {
        self.free_blocks.len()
    }

    /// Get the number of allocated blocks.
    pub fn num_allocated(&self) -> usize {
        self.total_blocks - self.free_blocks.len()
    }

    /// Get the block size.
    pub fn block_size(&self) -> usize {
        self.block_size
    }

    /// Get total blocks.
    pub fn total_blocks(&self) -> usize {
        self.total_blocks
    }
}

/// Block table mapping logical to physical blocks for a sequence.
#[derive(Debug, Clone)]
pub struct BlockTable {
    /// Logical to physical block mapping.
    block_indices: Vec<usize>,
    /// Number of tokens stored.
    num_tokens: usize,
    /// Block size.
    block_size: usize,
}

impl BlockTable {
    /// Create a new block table.
    pub fn new(block_size: usize) -> Self {
        Self {
            block_indices: Vec::new(),
            num_tokens: 0,
            block_size,
        }
    }

    /// Get the number of blocks.
    pub fn num_blocks(&self) -> usize {
        self.block_indices.len()
    }

    /// Get the number of tokens.
    pub fn num_tokens(&self) -> usize {
        self.num_tokens
    }

    /// Get the block indices.
    pub fn block_indices(&self) -> &[usize] {
        &self.block_indices
    }

    /// Add a block to the table.
    pub fn add_block(&mut self, block_idx: usize) {
        self.block_indices.push(block_idx);
    }

    /// Add tokens to the table, returning number of new blocks needed.
    pub fn add_tokens(&mut self, num_tokens: usize) -> usize {
        let old_blocks = (self.num_tokens + self.block_size - 1) / self.block_size;
        self.num_tokens += num_tokens;
        let new_blocks = (self.num_tokens + self.block_size - 1) / self.block_size;
        new_blocks.saturating_sub(old_blocks)
    }

    /// Get the physical block and offset for a token position.
    pub fn get_block_and_offset(&self, token_pos: usize) -> Option<(usize, usize)> {
        let block_idx = token_pos / self.block_size;
        let offset = token_pos % self.block_size;
        self.block_indices
            .get(block_idx)
            .map(|&phys| (phys, offset))
    }
}

/// Paged KV cache for efficient batched inference.
///
/// This cache uses block-based memory management for:
/// - Memory-efficient variable-length batching
/// - Near-zero memory fragmentation
/// - Efficient block reuse across sequences
///
/// # Architecture
///
/// ```text
/// ┌─────────────────────────────────────────────────────────────┐
/// │                    Physical Blocks                          │
/// │ [Block 0][Block 1][Block 2][Block 3][Block 4]...           │
/// │    K+V      K+V      K+V      K+V      K+V                  │
/// └─────────────────────────────────────────────────────────────┘
///                    ↑ ↑ ↑
/// ┌──────────────────┘ │ └──────────────────┐
/// │                    │                    │
/// │ Sequence 0        Sequence 1           Sequence 2          │
/// │ [0, 3]            [1, 4]               [2]                  │
/// │ (64 tokens)       (64 tokens)          (32 tokens)          │
/// └─────────────────────────────────────────────────────────────┘
/// ```
#[derive(Debug)]
pub struct PagedKVCache {
    /// Configuration.
    config: PagedKVCacheConfig,
    /// Block allocator.
    allocator: BlockAllocator,
    /// Physical key blocks per layer [layer][block][kv_heads, block_size, head_dim].
    key_blocks: Vec<Vec<Option<Array>>>,
    /// Physical value blocks per layer [layer][block][kv_heads, block_size, head_dim].
    value_blocks: Vec<Vec<Option<Array>>>,
    /// Block tables per sequence.
    block_tables: std::collections::HashMap<u64, BlockTable>,
    /// Next sequence ID.
    next_seq_id: u64,
}

impl PagedKVCache {
    /// Create a new paged KV cache.
    pub fn new(config: PagedKVCacheConfig) -> Self {
        let num_layers = config.num_layers;
        let max_blocks = config.max_blocks;

        // Pre-allocate block storage (but not the actual arrays yet - lazy allocation)
        let key_blocks: Vec<Vec<Option<Array>>> =
            (0..num_layers).map(|_| vec![None; max_blocks]).collect();
        let value_blocks: Vec<Vec<Option<Array>>> =
            (0..num_layers).map(|_| vec![None; max_blocks]).collect();

        Self {
            allocator: BlockAllocator::new(max_blocks, config.block_size),
            key_blocks,
            value_blocks,
            block_tables: std::collections::HashMap::new(),
            next_seq_id: 0,
            config,
        }
    }

    /// Allocate a new sequence, returning its ID.
    ///
    /// # Arguments
    /// * `initial_tokens` - Number of tokens to allocate initially (typically prompt length)
    pub fn allocate_sequence(&mut self, initial_tokens: usize) -> Result<u64, Exception> {
        let num_blocks = (initial_tokens + self.config.block_size - 1) / self.config.block_size;
        let blocks = self
            .allocator
            .allocate_n(num_blocks)
            .ok_or_else(|| Exception::custom("Out of KV cache blocks"))?;

        let seq_id = self.next_seq_id;
        self.next_seq_id += 1;

        let mut table = BlockTable::new(self.config.block_size);
        for block_idx in blocks {
            table.add_block(block_idx);
            // Lazy allocation: initialize block arrays on first use
            self.ensure_block_allocated(block_idx)?;
        }
        table.num_tokens = initial_tokens;

        self.block_tables.insert(seq_id, table);
        Ok(seq_id)
    }

    /// Extend a sequence with additional tokens.
    pub fn extend_sequence(&mut self, seq_id: u64, num_tokens: usize) -> Result<(), Exception> {
        // Calculate new blocks needed without holding borrow
        let new_blocks_needed = {
            let table = self
                .block_tables
                .get_mut(&seq_id)
                .ok_or_else(|| Exception::custom("Sequence not found"))?;
            table.add_tokens(num_tokens)
        };

        // Allocate and ensure blocks
        let mut new_block_indices = Vec::new();
        for _ in 0..new_blocks_needed {
            let block_idx = self
                .allocator
                .allocate()
                .ok_or_else(|| Exception::custom("Out of KV cache blocks"))?;
            self.ensure_block_allocated(block_idx)?;
            new_block_indices.push(block_idx);
        }

        // Add blocks to table
        if let Some(table) = self.block_tables.get_mut(&seq_id) {
            for block_idx in new_block_indices {
                table.add_block(block_idx);
            }
        }

        Ok(())
    }

    /// Free a sequence and return its blocks.
    pub fn free_sequence(&mut self, seq_id: u64) {
        if let Some(table) = self.block_tables.remove(&seq_id) {
            self.allocator.free_all(table.block_indices());
        }
    }

    /// Update KV cache for a sequence at a specific layer.
    ///
    /// # Arguments
    /// * `seq_id` - Sequence ID
    /// * `layer_idx` - Layer index
    /// * `new_keys` - New keys [batch=1, kv_heads, new_seq, head_dim]
    /// * `new_values` - New values [batch=1, kv_heads, new_seq, head_dim]
    /// * `start_pos` - Starting position in the sequence
    pub fn update(
        &mut self,
        seq_id: u64,
        layer_idx: usize,
        new_keys: &Array,
        new_values: &Array,
        start_pos: usize,
    ) -> Result<(), Exception> {
        let num_new_tokens = new_keys.dim(2) as usize;

        // Collect all block/offset pairs first to avoid holding table borrow
        let block_offsets: Vec<(usize, usize)> = {
            let table = self
                .block_tables
                .get(&seq_id)
                .ok_or_else(|| Exception::custom("Sequence not found"))?;

            (0..num_new_tokens)
                .map(|i| {
                    let token_pos = start_pos + i;
                    table.get_block_and_offset(token_pos)
                })
                .collect::<Option<Vec<_>>>()
                .ok_or_else(|| Exception::custom("Token position out of range"))?
        };

        // Now update blocks
        for (i, (block_idx, offset)) in block_offsets.into_iter().enumerate() {
            // Slice the single token from input
            let k_token = new_keys.index((.., .., i as i32..=i as i32, ..));
            let v_token = new_values.index((.., .., i as i32..=i as i32, ..));

            // Update block at offset
            self.update_block_at_offset(layer_idx, block_idx, offset, &k_token, &v_token)?;
        }

        Ok(())
    }

    /// Fetch cached K/V for attention computation.
    ///
    /// Returns concatenated K/V arrays for all tokens in the sequence.
    pub fn fetch(&self, seq_id: u64, layer_idx: usize) -> Result<(Array, Array), Exception> {
        let table = self
            .block_tables
            .get(&seq_id)
            .ok_or_else(|| Exception::custom("Sequence not found"))?;

        let num_tokens = table.num_tokens();
        if num_tokens == 0 {
            return Err(Exception::custom("Empty sequence"));
        }

        // Gather blocks and concatenate
        let mut key_parts: Vec<Array> = Vec::new();
        let mut value_parts: Vec<Array> = Vec::new();

        let block_size = self.config.block_size;
        let mut remaining = num_tokens;

        for &block_idx in table.block_indices().iter() {
            let tokens_in_block = remaining.min(block_size);

            if let (Some(k_block), Some(v_block)) = (
                &self.key_blocks[layer_idx][block_idx],
                &self.value_blocks[layer_idx][block_idx],
            ) {
                // Slice the valid portion of the block
                let k_slice = if tokens_in_block < block_size {
                    k_block.index((.., ..tokens_in_block as i32, ..))
                } else {
                    k_block.clone()
                };
                let v_slice = if tokens_in_block < block_size {
                    v_block.index((.., ..tokens_in_block as i32, ..))
                } else {
                    v_block.clone()
                };

                key_parts.push(k_slice);
                value_parts.push(v_slice);
            }

            remaining -= tokens_in_block;
            if remaining == 0 {
                break;
            }
        }

        // Concatenate all blocks along sequence dimension
        if key_parts.is_empty() {
            return Err(Exception::custom("No blocks to fetch"));
        }

        let key_refs: Vec<&Array> = key_parts.iter().collect();
        let value_refs: Vec<&Array> = value_parts.iter().collect();

        // Reshape to [batch=1, heads, seq, dim] format expected by attention
        let keys = concatenate_axis(&key_refs, 1)?;
        let values = concatenate_axis(&value_refs, 1)?;

        // Reshape from [heads, seq, dim] to [1, heads, seq, dim]
        let keys = keys.expand_dims(0)?;
        let values = values.expand_dims(0)?;

        Ok((keys, values))
    }

    /// Get the block table for a sequence (for kernel dispatch).
    pub fn get_block_table(&self, seq_id: u64) -> Option<&BlockTable> {
        self.block_tables.get(&seq_id)
    }

    /// Get number of sequences.
    pub fn num_sequences(&self) -> usize {
        self.block_tables.len()
    }

    /// Get memory statistics.
    pub fn memory_stats(&self) -> PagedCacheMemoryStats {
        let block_elements =
            self.config.num_kv_heads * self.config.block_size * self.config.head_dim;
        let bytes_per_block = block_elements * dtype_size(self.config.dtype) * 2; // K + V

        PagedCacheMemoryStats {
            total_blocks: self.allocator.total_blocks(),
            allocated_blocks: self.allocator.num_allocated(),
            free_blocks: self.allocator.num_free(),
            bytes_per_block,
            total_memory_bytes: self.allocator.total_blocks()
                * bytes_per_block
                * self.config.num_layers,
            used_memory_bytes: self.allocator.num_allocated()
                * bytes_per_block
                * self.config.num_layers,
        }
    }

    /// Reset the cache, freeing all sequences.
    pub fn reset(&mut self) {
        // Free all block tables
        let seq_ids: Vec<u64> = self.block_tables.keys().cloned().collect();
        for seq_id in seq_ids {
            self.free_sequence(seq_id);
        }
        self.next_seq_id = 0;
    }

    /// Ensure a block is allocated (lazy allocation).
    fn ensure_block_allocated(&mut self, block_idx: usize) -> Result<(), Exception> {
        let shape = [
            self.config.num_kv_heads as i32,
            self.config.block_size as i32,
            self.config.head_dim as i32,
        ];

        for layer_idx in 0..self.config.num_layers {
            if self.key_blocks[layer_idx][block_idx].is_none() {
                self.key_blocks[layer_idx][block_idx] = Some(Array::zeros::<f32>(&shape)?);
                self.value_blocks[layer_idx][block_idx] = Some(Array::zeros::<f32>(&shape)?);
            }
        }
        Ok(())
    }

    /// Update a block at a specific offset.
    fn update_block_at_offset(
        &mut self,
        layer_idx: usize,
        block_idx: usize,
        offset: usize,
        key: &Array,
        value: &Array,
    ) -> Result<(), Exception> {
        // Get or create the block as mutable
        let k_block = self.key_blocks[layer_idx][block_idx]
            .as_mut()
            .ok_or_else(|| Exception::custom("Block not allocated"))?;
        let v_block = self.value_blocks[layer_idx][block_idx]
            .as_mut()
            .ok_or_else(|| Exception::custom("Block not allocated"))?;

        // Remove batch dimension from input [1, heads, 1, dim] -> [heads, 1, dim]
        let k_squeezed = key.squeeze_axes(&[0])?;
        let v_squeezed = value.squeeze_axes(&[0])?;

        // In-place update using TryIndexMutOp (SOTA O(1) update)
        k_block.try_index_mut((.., offset as i32..=offset as i32, ..), &k_squeezed)?;
        v_block.try_index_mut((.., offset as i32..=offset as i32, ..), &v_squeezed)?;

        Ok(())
    }
}

/// Memory statistics for paged cache.
#[derive(Debug, Clone)]
pub struct PagedCacheMemoryStats {
    /// Total number of blocks.
    pub total_blocks: usize,
    /// Number of allocated blocks.
    pub allocated_blocks: usize,
    /// Number of free blocks.
    pub free_blocks: usize,
    /// Bytes per block.
    pub bytes_per_block: usize,
    /// Total memory in bytes.
    pub total_memory_bytes: usize,
    /// Used memory in bytes.
    pub used_memory_bytes: usize,
}

impl PagedCacheMemoryStats {
    /// Get memory utilization as a percentage.
    pub fn utilization(&self) -> f64 {
        if self.total_blocks == 0 {
            0.0
        } else {
            (self.allocated_blocks as f64 / self.total_blocks as f64) * 100.0
        }
    }
}

/// Convenience function to create a paged KV cache.
pub fn create_paged_cache(
    num_layers: usize,
    num_kv_heads: usize,
    head_dim: usize,
    max_seq_len: usize,
) -> PagedKVCache {
    PagedKVCache::new(PagedKVCacheConfig::new(
        num_layers,
        num_kv_heads,
        head_dim,
        max_seq_len,
    ))
}

// ============================================================================
// Mamba SSM State Cache
// ============================================================================

/// Cache for Mamba-2 SSM state during autoregressive generation.
///
/// Mamba layers require two types of state for incremental generation:
/// 1. **Conv state**: Last (kernel_size - 1) conv1d inputs for causal convolution
/// 2. **SSM state**: The hidden state matrix from the state space model
///
/// Without this cache, each generated token is processed without context from
/// previous tokens through Mamba layers, producing incoherent output.
#[derive(Debug, Clone)]
pub struct MambaCache {
    /// Per-layer cache entries.
    /// Each entry is (conv_state, ssm_state) where both may be None initially.
    layers: Vec<MambaCacheEntry>,
}

/// Cache entry for a single Mamba layer.
#[derive(Debug, Clone, Default)]
pub struct MambaCacheEntry {
    /// Convolutional state - last (kernel_size - 1) inputs.
    /// Shape: [batch, kernel_size - 1, conv_dim]
    pub conv_state: Option<Array>,
    /// SSM hidden state.
    /// Shape: [batch, num_heads, head_dim, state_dim]
    pub ssm_state: Option<Array>,
}

impl MambaCache {
    /// Create a new Mamba cache with the specified number of layers.
    pub fn new(num_layers: usize) -> Self {
        let layers = (0..num_layers)
            .map(|_| MambaCacheEntry::default())
            .collect();
        Self { layers }
    }

    /// Get a mutable reference to a layer's cache entry.
    pub fn get_mut(&mut self, layer_idx: usize) -> Option<&mut MambaCacheEntry> {
        self.layers.get_mut(layer_idx)
    }

    /// Get an immutable reference to a layer's cache entry.
    pub fn get(&self, layer_idx: usize) -> Option<&MambaCacheEntry> {
        self.layers.get(layer_idx)
    }

    /// Reset all cache entries to None.
    pub fn reset(&mut self) {
        for entry in &mut self.layers {
            entry.conv_state = None;
            entry.ssm_state = None;
        }
    }

    /// Check if the cache is empty (no state stored).
    pub fn is_empty(&self) -> bool {
        self.layers
            .iter()
            .all(|e| e.conv_state.is_none() && e.ssm_state.is_none())
    }
}

impl MambaCacheEntry {
    /// Update the conv state with new input, returning the padded input for conv1d.
    ///
    /// This implements causal convolution by:
    /// 1. Concatenating stored state with new input
    /// 2. Storing the last (kernel_size - 1) values for next call
    /// 3. Returning the padded input for conv1d processing
    ///
    /// # Arguments
    /// * `input` - New input tensor [batch, seq_len, conv_dim]
    /// * `kernel_size` - Conv1d kernel size
    ///
    /// # Returns
    /// Padded input [batch, seq_len + kernel_size - 1, conv_dim]
    pub fn update_conv_state(
        &mut self,
        input: &Array,
        kernel_size: i32,
    ) -> Result<Array, Exception> {
        let pad_len = (kernel_size - 1) as usize;
        let shape = input.shape();
        let batch = shape[0] as i32;
        let conv_dim = shape[2] as i32;

        // Get or initialize conv state with matching dtype
        let conv_state = if let Some(ref state) = self.conv_state {
            state.clone()
        } else {
            // Initialize to zeros with shape [batch, pad_len, conv_dim]
            // Use same dtype as input to avoid dtype mismatch issues
            Array::zeros::<f32>(&[batch, pad_len as i32, conv_dim])?.as_dtype(input.dtype())?
        };

        // Concatenate state with new input along sequence dimension
        let padded = concatenate_axis(&[&conv_state, input], 1)?;

        // Store last (kernel_size - 1) values for next call
        let seq_len = padded.dim(1);
        let start_idx = seq_len - pad_len as i32;
        self.conv_state = Some(padded.index((.., start_idx.., ..)));

        Ok(padded)
    }

    /// Get the current SSM state, returning None if not initialized.
    pub fn get_ssm_state(&self) -> Option<&Array> {
        self.ssm_state.as_ref()
    }

    /// Update the SSM state.
    pub fn set_ssm_state(&mut self, state: Array) {
        self.ssm_state = Some(state);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_kv_cache_config() {
        let config = KVCacheConfig::new(32, 2048, 8, 128)
            .with_dtype(Dtype::Float16)
            .with_sliding_window(512);

        assert_eq!(config.num_layers, 32);
        assert_eq!(config.max_seq_len, 2048);
        assert_eq!(config.num_kv_heads, 8);
        assert_eq!(config.head_dim, 128);
        assert_eq!(config.dtype, Dtype::Float16);
        assert_eq!(config.mode, CacheMode::SlidingWindow { window_size: 512 });
    }

    #[test]
    fn test_kv_cache_basic() {
        let config = KVCacheConfig::new(2, 100, 4, 64);
        let mut cache = KVCache::new(config);

        assert!(cache.is_empty());
        assert_eq!(cache.seq_len(), 0);

        // First update - [B, heads, seq, head_dim] format
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();

        let (cached_k, cached_v) = cache.update_and_fetch(0, &keys, &values).unwrap();

        // Seq is now axis 2
        assert_eq!(cached_k.dim(2), 10);
        assert_eq!(cached_v.dim(2), 10);
        assert_eq!(cache.seq_len(), 10);
        assert!(!cache.is_empty());
    }

    #[test]
    fn test_kv_cache_accumulation() {
        let config = KVCacheConfig::new(1, 100, 4, 64);
        let mut cache = KVCache::new(config);

        // First update: 10 tokens [B, heads, seq, head_dim]
        let k1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(0, &k1, &v1).unwrap();

        assert_eq!(cache.seq_len(), 10);

        // Second update: 5 more tokens
        let k2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let (cached_k, cached_v) = cache.update_and_fetch(0, &k2, &v2).unwrap();

        // Seq is axis 2
        assert_eq!(cached_k.dim(2), 15);
        assert_eq!(cached_v.dim(2), 15);
        assert_eq!(cache.seq_len(), 15);
        assert_eq!(cache.total_tokens(), 15);
    }

    #[test]
    fn test_kv_cache_sliding_window() {
        let config = KVCacheConfig::new(1, 100, 4, 64).with_sliding_window(20);
        let mut cache = KVCache::new(config);

        // Add 15 tokens [B, heads, seq, head_dim]
        let k1 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        cache.update_and_fetch(0, &k1, &v1).unwrap();

        assert_eq!(cache.seq_len(), 15);

        // Add 10 more - should trigger sliding window
        let k2 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let (cached_k, _) = cache.update_and_fetch(0, &k2, &v2).unwrap();

        // Should be trimmed to window size of 20, seq is axis 2
        assert_eq!(cached_k.dim(2), 20);
        assert_eq!(cache.seq_len(), 20);
        // But total tokens should reflect actual count
        assert_eq!(cache.total_tokens(), 25);
    }

    #[test]
    fn test_kv_cache_reset() {
        let config = KVCacheConfig::new(2, 100, 4, 64);
        let mut cache = KVCache::new(config);

        // [B, heads, seq, head_dim] format
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(0, &keys, &values).unwrap();
        cache.update_and_fetch(1, &keys, &values).unwrap();

        assert!(!cache.is_empty());

        cache.reset();

        assert!(cache.is_empty());
        assert_eq!(cache.seq_len(), 0);
        assert_eq!(cache.total_tokens(), 0);
    }

    #[test]
    fn test_kv_cache_rope_offset() {
        let config = KVCacheConfig::new(1, 100, 4, 64);
        let mut cache = KVCache::new(config);

        assert_eq!(cache.rope_offset(), 0);

        // [B, heads, seq, head_dim] format
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(0, &keys, &values).unwrap();

        assert_eq!(cache.rope_offset(), 10);
    }

    #[test]
    fn test_kv_cache_memory_estimation() {
        let config = KVCacheConfig::new(32, 2048, 8, 128).with_dtype(Dtype::Float16);
        let cache = KVCache::new(config);

        // 32 layers * 2048 seq * 8 heads * 128 dim * 2 bytes * 2 (K+V)
        let expected = 32 * 2048 * 8 * 128 * 2 * 2;
        assert_eq!(cache.max_memory_usage(), expected);
    }

    #[test]
    fn test_kv_cache_multi_layer() {
        let config = KVCacheConfig::new(4, 100, 4, 64);
        let mut cache = KVCache::new(config);

        // Update all layers - [B, heads, seq, head_dim] format
        for layer in 0..4 {
            let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
            let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
            cache.update_and_fetch(layer, &keys, &values).unwrap();
        }

        // All layers should have same seq_len (axis 2)
        for layer in 0..4 {
            let (k, v) = cache.get(layer).expect("Cache should exist");
            assert_eq!(k.dim(2), 10);
            assert_eq!(v.dim(2), 10);
        }
    }

    #[test]
    fn test_batch_kv_cache() {
        let config = KVCacheConfig::new(2, 100, 4, 64);
        let mut batch_cache = BatchKVCache::new(4, config);

        assert_eq!(batch_cache.batch_size(), 4);

        // Update one cache - [B, heads, seq, head_dim] format
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        batch_cache
            .get_mut(0)
            .unwrap()
            .update_and_fetch(0, &keys, &values)
            .unwrap();

        assert!(!batch_cache.caches[0].is_empty());
        assert!(batch_cache.caches[1].is_empty());
    }

    #[test]
    fn test_batch_kv_cache_reset_indices() {
        let config = KVCacheConfig::new(1, 100, 4, 64);
        let mut batch_cache = BatchKVCache::new(4, config);

        // [B, heads, seq, head_dim] format
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();

        // Fill all caches
        for i in 0..4 {
            batch_cache
                .get_mut(i)
                .unwrap()
                .update_and_fetch(0, &keys, &values)
                .unwrap();
        }

        // Reset specific indices
        batch_cache.reset_indices(&[0, 2]);

        assert!(batch_cache.caches[0].is_empty());
        assert!(!batch_cache.caches[1].is_empty());
        assert!(batch_cache.caches[2].is_empty());
        assert!(!batch_cache.caches[3].is_empty());
    }

    #[test]
    fn test_convenience_functions() {
        let cache = create_kv_cache(32, 2048, 8, 128);
        assert_eq!(cache.config().num_layers, 32);
        assert_eq!(cache.config().mode, CacheMode::Standard);

        let sliding_cache = create_sliding_window_cache(32, 512, 8, 128);
        assert_eq!(
            sliding_cache.config().mode,
            CacheMode::SlidingWindow { window_size: 512 }
        );
    }

    // =========================================================================
    // RotatingKVCache Tests
    // =========================================================================

    #[test]
    fn test_rotating_cache_basic() {
        let mut cache = RotatingKVCache::new(100, 0);

        assert!(cache.is_empty());
        assert_eq!(cache.len(), 0);
        assert_eq!(cache.offset(), 0);

        // First update - [B, heads, seq, head_dim] format
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();

        let (cached_k, cached_v) = cache.update_and_fetch(&keys, &values).unwrap();

        assert_eq!(cached_k.dim(2), 10);
        assert_eq!(cached_v.dim(2), 10);
        assert_eq!(cache.len(), 10);
        assert_eq!(cache.offset(), 10);
        assert!(!cache.is_empty());
    }

    #[test]
    fn test_rotating_cache_accumulation() {
        let mut cache = RotatingKVCache::new(100, 0);

        // First update: 10 tokens
        let k1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&k1, &v1).unwrap();

        assert_eq!(cache.len(), 10);

        // Second update: 5 more tokens
        let k2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let (cached_k, cached_v) = cache.update_and_fetch(&k2, &v2).unwrap();

        assert_eq!(cached_k.dim(2), 15);
        assert_eq!(cached_v.dim(2), 15);
        assert_eq!(cache.len(), 15);
    }

    #[test]
    fn test_rotating_cache_rotation() {
        let mut cache = RotatingKVCache::new(20, 0);

        // Fill beyond max_size
        let k1 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        cache.update_and_fetch(&k1, &v1).unwrap();

        let k2 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let (_cached_k, _) = cache.update_and_fetch(&k2, &v2).unwrap();

        // MLX-LM allows max_size + S - 1 to ensure every token gets at least max_size context
        // So cache can grow to max_size + num_steps - 1 before trimming
        // In this case: 20 + 10 - 1 = 29, then trimmed to ~20 region
        // The key behavior is len() is capped at max_size, offset tracks total
        assert!(cache.len() <= 25); // len() caps at max_size or slightly above during concat
        assert_eq!(cache.offset(), 25); // Total tokens seen
    }

    #[test]
    fn test_rotating_cache_with_keep() {
        let mut cache = RotatingKVCache::new(20, 4); // Keep first 4 tokens

        // Fill cache
        let k1 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        cache.update_and_fetch(&k1, &v1).unwrap();

        assert_eq!(cache.len(), 15);

        // Add more tokens
        let k2 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&k2, &v2).unwrap();

        // Cache should have rotated but kept initial tokens
        assert!(cache.len() <= 20);
    }

    #[test]
    fn test_rotating_cache_reset() {
        let mut cache = RotatingKVCache::new(100, 0);

        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&keys, &values).unwrap();

        assert!(!cache.is_empty());

        cache.reset();

        assert!(cache.is_empty());
        assert_eq!(cache.len(), 0);
        assert_eq!(cache.offset(), 0);
    }

    #[test]
    fn test_rotating_cache_single_token_updates() {
        let mut cache = RotatingKVCache::new(50, 0);

        // Simulate autoregressive generation
        for i in 0..30 {
            let k = Array::ones::<f32>(&[1, 4, 1, 64]).unwrap();
            let v = Array::ones::<f32>(&[1, 4, 1, 64]).unwrap();
            let (cached_k, _) = cache.update_and_fetch(&k, &v).unwrap();

            assert_eq!(cached_k.dim(2) as usize, (i + 1).min(50));
        }

        assert_eq!(cache.offset(), 30);
        assert_eq!(cache.len(), 30);
    }

    #[test]
    fn test_rotating_cache_trimmable() {
        let mut cache = RotatingKVCache::new(20, 0);

        assert!(cache.is_trimmable()); // Empty cache is trimmable

        let k = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&k, &v).unwrap();

        assert!(cache.is_trimmable()); // Under max_size

        // Fill to max
        let k2 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 15, 64]).unwrap();
        cache.update_and_fetch(&k2, &v2).unwrap();

        assert!(!cache.is_trimmable()); // At or over max_size
    }

    #[test]
    fn test_rotating_cache_rope_offset() {
        let mut cache = RotatingKVCache::new(100, 0);

        assert_eq!(cache.rope_offset(), 0);

        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&keys, &values).unwrap();

        // RoPE offset should be total tokens seen, not cache length
        assert_eq!(cache.rope_offset(), 10);
    }

    // =========================================================================
    // QuantizedKVCache Tests
    // =========================================================================

    #[test]
    fn test_quantized_cache_basic() {
        let mut cache = QuantizedKVCache::new(8, 64);

        assert!(cache.is_empty());
        assert_eq!(cache.len(), 0);

        // First update
        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();

        let (cached_k, cached_v) = cache.update_and_fetch(&keys, &values).unwrap();

        // Dequantized output should have correct shape
        assert_eq!(cached_k.dim(2), 10);
        assert_eq!(cached_v.dim(2), 10);
        assert_eq!(cache.len(), 10);
        assert!(!cache.is_empty());
    }

    #[test]
    fn test_quantized_cache_accumulation() {
        let mut cache = QuantizedKVCache::new(4, 64);

        // First update
        let k1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&k1, &v1).unwrap();

        // Second update
        let k2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let (cached_k, cached_v) = cache.update_and_fetch(&k2, &v2).unwrap();

        assert_eq!(cached_k.dim(2), 15);
        assert_eq!(cached_v.dim(2), 15);
        assert_eq!(cache.len(), 15);
    }

    #[test]
    fn test_quantized_cache_reset() {
        let mut cache = QuantizedKVCache::new(8, 64);

        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&keys, &values).unwrap();

        assert!(!cache.is_empty());

        cache.reset();

        assert!(cache.is_empty());
        assert_eq!(cache.len(), 0);
    }

    #[test]
    fn test_quantized_cache_different_bits() {
        // Test 8-bit
        let cache_8bit = QuantizedKVCache::new(8, 64);
        assert_eq!(cache_8bit.bits, 8);

        // Test 4-bit
        let cache_4bit = QuantizedKVCache::new(4, 64);
        assert_eq!(cache_4bit.bits, 4);

        // Test 2-bit
        let cache_2bit = QuantizedKVCache::new(2, 64);
        assert_eq!(cache_2bit.bits, 2);
    }

    #[test]
    #[should_panic(expected = "bits must be 2, 4, or 8")]
    fn test_quantized_cache_invalid_bits() {
        let _ = QuantizedKVCache::new(3, 64); // Invalid
    }

    #[test]
    fn test_quantized_cache_memory_usage() {
        let mut cache = QuantizedKVCache::new(8, 64);

        assert_eq!(cache.memory_usage(), 0); // Empty

        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&keys, &values).unwrap();

        // Should have some memory usage now
        assert!(cache.memory_usage() > 0);
    }

    #[test]
    fn test_quantized_cache_rope_offset() {
        let mut cache = QuantizedKVCache::new(8, 64);

        assert_eq!(cache.rope_offset(), 0);

        let keys = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 64]).unwrap();
        cache.update_and_fetch(&keys, &values).unwrap();

        assert_eq!(cache.rope_offset(), 10);
    }

    // =========================================================================
    // Convenience Function Tests for New Cache Types
    // =========================================================================

    #[test]
    fn test_create_rotating_cache() {
        let cache = create_rotating_cache(1024, 4);
        assert_eq!(cache.max_size, 1024);
        assert_eq!(cache.keep, 4);
        assert!(cache.is_empty());
    }

    #[test]
    fn test_create_quantized_cache() {
        let cache = create_quantized_cache(4, 32);
        assert_eq!(cache.bits, 4);
        assert_eq!(cache.group_size, 32);
        assert!(cache.is_empty());
    }

    // =========================================================================
    // CacheMode Tests for New Modes
    // =========================================================================

    #[test]
    fn test_cache_mode_rotating() {
        let config = KVCacheConfig::new(32, 2048, 8, 128).with_rotating(1024, 8);

        assert_eq!(
            config.mode,
            CacheMode::Rotating {
                max_size: 1024,
                keep: 8
            }
        );
    }

    #[test]
    fn test_cache_mode_quantized() {
        let config = KVCacheConfig::new(32, 2048, 8, 128).with_quantized(4, 64);

        assert_eq!(
            config.mode,
            CacheMode::Quantized {
                bits: 4,
                group_size: 64
            }
        );
    }

    // =========================================================================
    // PagedKVCache Tests
    // =========================================================================

    #[test]
    fn test_paged_cache_config() {
        let config = PagedKVCacheConfig::new(32, 8, 128, 2048)
            .with_block_size(16)
            .with_dtype(Dtype::Float16);

        assert_eq!(config.num_layers, 32);
        assert_eq!(config.num_kv_heads, 8);
        assert_eq!(config.head_dim, 128);
        assert_eq!(config.block_size, 16);
        assert_eq!(config.dtype, Dtype::Float16);
    }

    #[test]
    fn test_block_allocator_basic() {
        let mut allocator = BlockAllocator::new(10, 32);

        assert_eq!(allocator.total_blocks(), 10);
        assert_eq!(allocator.num_free(), 10);
        assert_eq!(allocator.num_allocated(), 0);

        // Allocate a block
        let block = allocator.allocate().unwrap();
        assert_eq!(allocator.num_allocated(), 1);
        assert_eq!(allocator.num_free(), 9);

        // Free the block
        allocator.free(block);
        assert_eq!(allocator.num_free(), 10);
        assert_eq!(allocator.num_allocated(), 0);
    }

    #[test]
    fn test_block_allocator_batch() {
        let mut allocator = BlockAllocator::new(10, 32);

        // Allocate 5 blocks at once
        let blocks = allocator.allocate_n(5).unwrap();
        assert_eq!(blocks.len(), 5);
        assert_eq!(allocator.num_allocated(), 5);

        // Try to allocate more than available
        assert!(allocator.allocate_n(6).is_none());

        // Free all
        allocator.free_all(&blocks);
        assert_eq!(allocator.num_free(), 10);
    }

    #[test]
    fn test_block_table() {
        let mut table = BlockTable::new(32);

        assert_eq!(table.num_tokens(), 0);
        assert_eq!(table.num_blocks(), 0);

        // Add a block and tokens
        table.add_block(5);
        table.num_tokens = 16;

        assert_eq!(table.num_tokens(), 16);
        assert_eq!(table.num_blocks(), 1);

        // Check block lookup
        let (phys, offset) = table.get_block_and_offset(10).unwrap();
        assert_eq!(phys, 5);
        assert_eq!(offset, 10);
    }

    #[test]
    fn test_block_table_add_tokens() {
        let mut table = BlockTable::new(32);
        table.add_block(0);

        // First 32 tokens need 1 block (tokens 0-31)
        let new_blocks = table.add_tokens(32);
        assert_eq!(new_blocks, 1); // Need 1 block for first 32 tokens

        // Adding 32 more requires 1 more block (tokens 32-63)
        table.add_block(1);
        let new_blocks = table.add_tokens(32);
        assert_eq!(new_blocks, 1);

        // Adding 1 more requires 1 new block (token 64)
        table.add_block(2);
        let new_blocks = table.add_tokens(1);
        assert_eq!(new_blocks, 1);
    }

    #[test]
    fn test_paged_cache_basic() {
        let config = PagedKVCacheConfig::new(2, 4, 64, 256);
        let mut cache = PagedKVCache::new(config);

        assert_eq!(cache.num_sequences(), 0);

        // Allocate a sequence
        let seq_id = cache.allocate_sequence(32).unwrap();
        assert_eq!(cache.num_sequences(), 1);

        // Check memory stats
        let stats = cache.memory_stats();
        assert!(stats.allocated_blocks > 0);
        assert!(stats.utilization() > 0.0);

        // Free the sequence
        cache.free_sequence(seq_id);
        assert_eq!(cache.num_sequences(), 0);
    }

    #[test]
    fn test_paged_cache_extend() {
        let config = PagedKVCacheConfig::new(1, 4, 64, 256);
        let mut cache = PagedKVCache::new(config);

        let seq_id = cache.allocate_sequence(16).unwrap();
        let table = cache.get_block_table(seq_id).unwrap();
        assert_eq!(table.num_tokens(), 16);

        // Extend the sequence
        cache.extend_sequence(seq_id, 20).unwrap();
        let table = cache.get_block_table(seq_id).unwrap();
        assert_eq!(table.num_tokens(), 36);
    }

    #[test]
    fn test_paged_cache_multiple_sequences() {
        let config = PagedKVCacheConfig::new(1, 4, 64, 1024);
        let mut cache = PagedKVCache::new(config);

        // Allocate multiple sequences
        let _seq1 = cache.allocate_sequence(64).unwrap();
        let seq2 = cache.allocate_sequence(32).unwrap();
        let _seq3 = cache.allocate_sequence(48).unwrap();

        assert_eq!(cache.num_sequences(), 3);

        // Free one
        cache.free_sequence(seq2);
        assert_eq!(cache.num_sequences(), 2);

        // Allocate new one (should reuse freed blocks)
        let _seq4 = cache.allocate_sequence(32).unwrap();
        assert_eq!(cache.num_sequences(), 3);

        // Reset all
        cache.reset();
        assert_eq!(cache.num_sequences(), 0);
    }

    #[test]
    fn test_paged_cache_memory_stats() {
        let config = PagedKVCacheConfig::new(2, 8, 128, 1024).with_dtype(Dtype::Float16);
        let mut cache = PagedKVCache::new(config);

        let stats = cache.memory_stats();
        assert_eq!(stats.allocated_blocks, 0);
        assert_eq!(stats.utilization(), 0.0);

        // Allocate a sequence
        let _ = cache.allocate_sequence(100);

        let stats = cache.memory_stats();
        assert!(stats.allocated_blocks > 0);
        assert!(stats.used_memory_bytes > 0);
    }

    #[test]
    fn test_create_paged_cache_convenience() {
        let cache = create_paged_cache(32, 8, 128, 2048);

        assert_eq!(cache.config.num_layers, 32);
        assert_eq!(cache.config.num_kv_heads, 8);
        assert_eq!(cache.config.head_dim, 128);
        assert_eq!(cache.num_sequences(), 0);
    }

    // =========================================================================
    // Eager Pre-Allocation Tests
    // =========================================================================

    #[test]
    fn test_kv_cache_eager_config() {
        let config = KVCacheConfig::new(32, 4096, 8, 128)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float16);

        assert!(config.eager_allocate);
        assert_eq!(config.eager_batch_size, 1);
        assert_eq!(config.max_seq_len, 4096);
        assert_eq!(config.dtype, Dtype::Float16);
    }

    #[test]
    fn test_kv_cache_eager_config_batch_size() {
        let config = KVCacheConfig::new(32, 2048, 8, 128).with_eager_allocate(4);

        assert!(config.eager_allocate);
        assert_eq!(config.eager_batch_size, 4);
    }

    #[test]
    fn test_kv_cache_memory_footprint() {
        // 32 layers × 2 (K+V) × 1 batch × 8 heads × 2048 seq × 128 dim × 2 bytes (fp16)
        let config = KVCacheConfig::new(32, 2048, 8, 128)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float16);

        let expected = 32 * 2 * 8 * 2048 * 128 * 2; // layers * K+V * heads * seq * dim * sizeof(f16)
        assert_eq!(config.memory_footprint(), expected);
    }

    #[test]
    fn test_kv_cache_memory_footprint_fp32() {
        // 16 layers × 2 (K+V) × 1 batch × 4 heads × 1024 seq × 64 dim × 4 bytes (fp32)
        let config = KVCacheConfig::new(16, 1024, 4, 64)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float32);

        let expected = 16 * 2 * 4 * 1024 * 64 * 4; // layers * K+V * heads * seq * dim * sizeof(f32)
        assert_eq!(config.memory_footprint(), expected);
    }

    #[test]
    fn test_kv_cache_memory_footprint_batch() {
        // Test with batch_size > 1
        let config = KVCacheConfig::new(8, 512, 4, 64)
            .with_eager_allocate(4)
            .with_dtype(Dtype::Float16);

        let expected = 8 * 2 * 4 * 4 * 512 * 64 * 2;
        assert_eq!(config.memory_footprint(), expected);
    }

    #[test]
    fn test_kv_cache_memory_footprint_human_bytes() {
        let config = KVCacheConfig::new(1, 1, 1, 1)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float32);

        // 1 × 2 × 1 × 1 × 1 × 1 × 4 = 8 bytes
        assert_eq!(config.memory_footprint_human(), "8 bytes");
    }

    #[test]
    fn test_kv_cache_memory_footprint_human_kb() {
        let config = KVCacheConfig::new(1, 128, 1, 1)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float32);

        // 1 × 2 × 1 × 1 × 128 × 1 × 4 = 1024 bytes = 1 KB
        assert_eq!(config.memory_footprint_human(), "1.00 KB");
    }

    #[test]
    fn test_kv_cache_memory_footprint_human_mb() {
        let config = KVCacheConfig::new(1, 2048, 8, 64)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float32);

        // 1 × 2 × 1 × 8 × 2048 × 64 × 4 = 8,388,608 bytes = 8 MB
        let human = config.memory_footprint_human();
        assert!(human.contains("MB"), "Expected MB, got: {}", human);
    }

    #[test]
    fn test_kv_cache_memory_footprint_human_gb() {
        // Need to exceed 1 GB (1,073,741,824 bytes)
        // 32 layers × 2 × 1 batch × 8 heads × 8192 seq × 128 dim × 2 bytes = 1 GB
        let config = KVCacheConfig::new(32, 8192, 8, 128)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float16);

        let human = config.memory_footprint_human();
        assert!(human.contains("GB"), "Expected GB, got: {}", human);
    }

    #[test]
    fn test_kv_cache_new_eager() {
        let config = KVCacheConfig::new(2, 128, 4, 64)
            .with_eager_allocate(1)
            .with_dtype(Dtype::Float32);

        let cache = KVCache::new_eager(config).expect("Should create eager cache");

        assert!(cache.is_preallocated());
        assert!(cache.is_empty()); // Pre-allocated but no data
        assert_eq!(cache.seq_len(), 0);
    }

    #[test]
    fn test_kv_cache_new_eager_fallback() {
        // When eager_allocate is false, new_eager should fall back to lazy
        let config = KVCacheConfig::new(2, 128, 4, 64);

        let cache = KVCache::new_eager(config).expect("Should create lazy cache");

        assert!(!cache.is_preallocated());
    }

    #[test]
    fn test_kv_cache_eager_is_preallocated() {
        // Eager cache should report as preallocated
        let eager_config = KVCacheConfig::new(2, 64, 4, 32).with_eager_allocate(1);
        let eager_cache = KVCache::new_eager(eager_config).unwrap();
        assert!(eager_cache.is_preallocated());

        // Lazy cache should not report as preallocated
        let lazy_config = KVCacheConfig::new(2, 64, 4, 32);
        let lazy_cache = KVCache::new(lazy_config);
        assert!(!lazy_cache.is_preallocated());
    }

    #[test]
    fn test_kv_cache_eager_is_empty() {
        let config = KVCacheConfig::new(2, 64, 4, 32).with_eager_allocate(1);
        let mut cache = KVCache::new_eager(config).unwrap();

        // Pre-allocated but empty (no data added)
        assert!(cache.is_empty());

        // Add some data
        let keys = Array::zeros::<f32>(&[1, 4, 10, 32]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 32]).unwrap();
        cache.update_and_fetch(0, &keys, &values).unwrap();

        // Now not empty
        assert!(!cache.is_empty());
    }

    #[test]
    fn test_kv_cache_eager_update() {
        let config = KVCacheConfig::new(2, 128, 4, 64).with_eager_allocate(1);
        let mut cache = KVCache::new_eager(config).unwrap();

        // First update
        let k1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 10, 64]).unwrap();
        let (cached_k, cached_v) = cache.update_and_fetch(0, &k1, &v1).unwrap();

        assert_eq!(cached_k.dim(2), 10);
        assert_eq!(cached_v.dim(2), 10);
        assert_eq!(cache.seq_len(), 10);

        // Second update (accumulation)
        let k2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 5, 64]).unwrap();
        let (cached_k, cached_v) = cache.update_and_fetch(0, &k2, &v2).unwrap();

        assert_eq!(cached_k.dim(2), 15);
        assert_eq!(cached_v.dim(2), 15);
        assert_eq!(cache.seq_len(), 15);
    }

    #[test]
    fn test_kv_cache_eager_reset_preserves_buffers() {
        let config = KVCacheConfig::new(2, 64, 4, 32).with_eager_allocate(1);
        let mut cache = KVCache::new_eager(config).unwrap();

        // Add some data
        let keys = Array::zeros::<f32>(&[1, 4, 10, 32]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 32]).unwrap();
        cache.update_and_fetch(0, &keys, &values).unwrap();
        assert!(!cache.is_empty());

        // Reset should preserve buffers in eager mode
        cache.reset();

        assert!(cache.is_empty()); // Offset reset
        assert!(cache.is_preallocated()); // Buffers preserved
        assert_eq!(cache.seq_len(), 0);
    }

    #[test]
    fn test_kv_cache_eager_reset_full_deallocates() {
        let config = KVCacheConfig::new(2, 64, 4, 32).with_eager_allocate(1);
        let mut cache = KVCache::new_eager(config).unwrap();

        assert!(cache.is_preallocated());

        // reset_full should deallocate even in eager mode
        cache.reset_full();

        assert!(!cache.is_preallocated()); // Buffers deallocated
        assert!(cache.is_empty());
    }

    #[test]
    fn test_kv_cache_lazy_reset_deallocates() {
        let config = KVCacheConfig::new(2, 64, 4, 32);
        let mut cache = KVCache::new(config);

        // Add some data
        let keys = Array::zeros::<f32>(&[1, 4, 10, 32]).unwrap();
        let values = Array::zeros::<f32>(&[1, 4, 10, 32]).unwrap();
        cache.update_and_fetch(0, &keys, &values).unwrap();
        assert!(!cache.is_empty());

        // Reset should deallocate in lazy mode
        cache.reset();

        assert!(cache.is_empty());
        // Verify buffers are deallocated by checking get returns None
        assert!(cache.get(0).is_none());
    }

    #[test]
    fn test_kv_cache_eager_reuse_after_reset() {
        let config = KVCacheConfig::new(2, 64, 4, 32).with_eager_allocate(1);
        let mut cache = KVCache::new_eager(config).unwrap();

        // First generation
        let k1 = Array::ones::<f32>(&[1, 4, 20, 32]).unwrap();
        let v1 = Array::ones::<f32>(&[1, 4, 20, 32]).unwrap();
        cache.update_and_fetch(0, &k1, &v1).unwrap();
        assert_eq!(cache.seq_len(), 20);

        // Reset for new generation
        cache.reset();
        assert!(cache.is_empty());
        assert!(cache.is_preallocated());

        // Second generation (reuses pre-allocated buffers)
        let k2 = Array::ones::<f32>(&[1, 4, 15, 32]).unwrap();
        let v2 = Array::ones::<f32>(&[1, 4, 15, 32]).unwrap();
        cache.update_and_fetch(0, &k2, &v2).unwrap();
        assert_eq!(cache.seq_len(), 15);
    }
}
