use cudarc::driver::{CudaDevice, CudaSlice};
use half::f16;
use std::sync::Arc;

#[derive(Clone)]
pub struct AttentionMemory {
    pub device: Arc<CudaDevice>,
    pub k_cache: CudaSlice<f16>,
    pub v_cache: CudaSlice<f16>,
    pub max_seq_len: usize,
    pub num_layers: usize,
    pub num_heads: usize,
    pub head_dim: usize,
    pub current_len: usize,
}

impl AttentionMemory {
    pub fn new(device: Arc<CudaDevice>, num_layers: usize, num_heads: usize, head_dim: usize, max_seq_len: usize) -> anyhow::Result<Self> {
        let elems = num_layers * num_heads * max_seq_len * head_dim;
        let k_cache = unsafe { device.alloc::<f16>(elems)? };
        let v_cache = unsafe { device.alloc::<f16>(elems)? };
        Ok(Self { device, k_cache, v_cache, max_seq_len, num_layers, num_heads, head_dim, current_len: 0 })
    }
    pub fn append_kv(&mut self, _layer: usize, _k: &CudaSlice<f16>, _v: &CudaSlice<f16>, seq_len: usize) -> anyhow::Result<()> {
        if self.current_len + seq_len > self.max_seq_len { anyhow::bail!("KV cache overflow"); }
        self.current_len += seq_len;
        Ok(())
    }
    pub fn reset(&mut self) { self.current_len = 0; }
}