KV cache implementation

pub trait KeyValueCache {
    /// Returns the current offset (number of tokens in cache)
    fn offset(&self) -> i32;

    /// Returns the maximum cache size (for sliding window), if any
    fn max_size(&self) -> Option<i32>;

    /// Update cache with new keys/values and return full cache contents
    fn update_and_fetch(&mut self, keys: Array, values: Array) 
        -> Result<(Array, Array), Exception>;
}

/// Simple concatenation-based KV cache
#[derive(Debug, Clone, Default)]
pub struct ConcatKeyValueCache {
    keys: Option<Array>,
    values: Option<Array>,
    offset: i32,
}

impl ConcatKeyValueCache {
    pub fn new() -> Self {
        Self::default()
    }
}

impl KeyValueCache for ConcatKeyValueCache {
    fn update_and_fetch(&mut self, keys: Array, values: Array) 
        -> Result<(Array, Array), Exception> {
        match (self.keys.take(), self.values.take()) {
            (Some(k), Some(v)) => {
                // Concatenate along sequence dimension (axis=-2)
                self.keys = Some(concatenate_axis(&[k, keys], -2)?);
                self.values = Some(concatenate_axis(&[v, values], -2)?);
            }
            _ => {
                // First token: just store
                self.keys = Some(keys);
                self.values = Some(values);
            }
        }
        
        let shape = self.keys.as_ref().expect("Keys cannot be None").shape();
        self.offset = shape[shape.len() - 2];

        Ok((
            self.keys.clone().expect("Keys cannot be None"),
            self.values.clone().expect("Values cannot be None"),
        ))
    }
}

/// Step-based KV Cache with pre-allocation (matches Python mlx-lm KVCache)
///
/// This cache pre-allocates buffers in steps of 256 tokens and uses in-place
/// slice updates, avoiding expensive concatenation on every token.
#[derive(Debug, Clone)]
pub struct KVCache {
    keys: Option<Array>,
    values: Option<Array>,
    offset: i32,
    step: i32,  // Default: 256
}

impl KVCache {
    pub fn new() -> Self {
        Self::with_step(256)
    }

    pub fn with_step(step: i32) -> Self {
        Self {
            keys: None,
            values: None,
            offset: 0,
            step,
        }
    }
}

fn update_and_fetch(&mut self, keys: Array, values: Array) 
    -> Result<(Array, Array), Exception> {
    let prev = self.offset;
    let num_new = keys.shape()[2];

    // Check if we need to grow the buffer
    let needs_grow = match &self.keys {
        None => true,
        Some(k) => (prev + num_new) > k.shape()[2],
    };

    if needs_grow {
        // Pre-allocate in multiples of step size (256)
        let n_steps = (self.step + num_new - 1) / self.step;
        let new_size = n_steps * self.step;
        
        let new_k = zeros_dtype(k_shape, k_dtype)?;
        let new_v = zeros_dtype(v_shape, v_dtype)?;
        
        // Concatenate old and new buffers
        self.keys = Some(concatenate_axis(&[old_k, new_k], 2)?);
        self.values = Some(concatenate_axis(&[old_v, new_v], 2)?);
    }

    self.offset += num_new;

    // In-place update using slice indexing
    let k = self.keys.as_mut().unwrap();
    let v = self.values.as_mut().unwrap();
    k.index_mut((Ellipsis, prev..self.offset, ..), &keys);
    v.index_mut((Ellipsis, prev..self.offset, ..), &values);

    Ok((
        k.index((Ellipsis, ..self.offset, ..)),
        v.index((Ellipsis, ..self.offset, ..)),
    ))
}