//! Simple pre-allocated KV cache. PagedAttention deferred to F2.

/// Per-layer KV cache buffer.
pub struct KvCache {
    pub k: Vec<f32>,
    pub v: Vec<f32>,
    pub len: usize,
    pub kv_dim: usize,
    pub max_seq_len: usize,
}

impl KvCache {
    pub fn new(max_seq_len: usize, num_kv_heads: usize, head_dim: usize) -> Self {
        let kv_dim = num_kv_heads * head_dim;
        Self {
            k: vec![0.0; max_seq_len * kv_dim],
            v: vec![0.0; max_seq_len * kv_dim],
            len: 0,
            kv_dim,
            max_seq_len,
        }
    }

    pub fn append(&mut self, new_k: &[f32], new_v: &[f32], n: usize) {
        let offset = self.len * self.kv_dim;
        let size = n * self.kv_dim;
        self.k[offset..offset + size].copy_from_slice(&new_k[..size]);
        self.v[offset..offset + size].copy_from_slice(&new_v[..size]);
        self.len += n;
    }

    pub fn k_slice(&self) -> &[f32] {
        &self.k[..self.len * self.kv_dim]
    }

    pub fn v_slice(&self) -> &[f32] {
        &self.v[..self.len * self.kv_dim]
    }

    pub fn reset(&mut self) {
        self.len = 0;
    }

    /// Total bytes allocated for this cache layer (K + V buffers).
    pub fn size_bytes(&self) -> usize {
        (self.k.len() + self.v.len()) * std::mem::size_of::<f32>()
    }
}

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

    #[test]
    fn test_kv_cache_append_and_retrieve() {
        let mut cache = KvCache::new(16, 2, 4);
        assert_eq!(cache.len, 0);
        let k = vec![1.0; 2 * 8];
        let v = vec![2.0; 2 * 8];
        cache.append(&k, &v, 2);
        assert_eq!(cache.len, 2);
        assert_eq!(cache.k_slice().len(), 16);
        assert!(cache.k_slice().iter().all(|&x| x == 1.0));
    }

    #[test]
    fn test_kv_cache_reset() {
        let mut cache = KvCache::new(16, 2, 4);
        cache.append(&[1.0; 8], &[2.0; 8], 1);
        assert_eq!(cache.len, 1);
        cache.reset();
        assert_eq!(cache.len, 0);
    }

    #[test]
    fn test_kv_cache_size_bytes() {
        let cache = KvCache::new(16, 2, 4);
        // kv_dim = 2 * 4 = 8
        // K buffer: 16 * 8 = 128 f32s = 512 bytes
        // V buffer: 16 * 8 = 128 f32s = 512 bytes
        // Total: 1024 bytes
        assert_eq!(cache.size_bytes(), 1024);
    }
}
