//! VRAM allocation tracker for LLMHosts Core.
//!
//! Tracks bytes allocated for model weights, KV caches, and scratch buffers.
//! Does NOT perform actual GPU memory allocation — CubeCL handles that.
//! This is bookkeeping so we can detect OOM before it happens and report
//! usage to the user.

use std::fmt;
use thiserror::Error;

/// Error returned when a memory allocation would exceed available VRAM.
#[derive(Error, Debug, Clone, PartialEq, Eq)]
pub enum MemoryError {
    #[error(
        "Out of memory: requested {requested} bytes but only {available} bytes available \
         (total: {total}, allocated: {allocated})"
    )]
    OutOfMemory {
        requested: usize,
        available: usize,
        total: usize,
        allocated: usize,
    },
}

/// Snapshot of memory usage at a point in time.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MemoryUsage {
    /// Bytes currently allocated.
    pub allocated: usize,
    /// Total bytes available to the allocator.
    pub total: usize,
    /// Bytes still free.
    pub free: usize,
    /// Bytes allocated for model weights.
    pub weights: usize,
    /// Bytes allocated for KV caches.
    pub kv_cache: usize,
}

impl fmt::Display for MemoryUsage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{:.1} / {:.1} MB used ({:.1} MB weights, {:.1} MB KV cache, {:.1} MB free)",
            self.allocated as f64 / 1_048_576.0,
            self.total as f64 / 1_048_576.0,
            self.weights as f64 / 1_048_576.0,
            self.kv_cache as f64 / 1_048_576.0,
            self.free as f64 / 1_048_576.0,
        )
    }
}

/// Tracks VRAM usage across model weights and KV caches.
///
/// This is a bookkeeping layer — it does not perform actual GPU memory
/// operations. It tracks how many bytes have been logically reserved so
/// we can prevent OOM and report usage.
pub struct GpuAllocator {
    /// Name of the backend this allocator tracks (e.g. "cuda:0", "wgpu:Metal").
    backend_name: String,
    /// Total VRAM budget in bytes.
    total_bytes: usize,
    /// Bytes allocated for model weights.
    weight_bytes: usize,
    /// Bytes allocated for KV caches.
    kv_cache_bytes: usize,
}

impl GpuAllocator {
    /// Create a new allocator with a given VRAM budget.
    pub fn new(backend_name: &str, total_bytes: usize) -> Self {
        Self {
            backend_name: backend_name.to_string(),
            total_bytes,
            weight_bytes: 0,
            kv_cache_bytes: 0,
        }
    }

    /// Backend name this allocator is tracking.
    pub fn backend_name(&self) -> &str {
        &self.backend_name
    }

    /// Total VRAM budget in bytes.
    pub fn total_bytes(&self) -> usize {
        self.total_bytes
    }

    /// Total bytes currently allocated (weights + KV cache).
    fn allocated(&self) -> usize {
        self.weight_bytes + self.kv_cache_bytes
    }

    /// Reserve bytes for model weights. Returns error if it would exceed budget.
    pub fn allocate_weights(&mut self, bytes: usize) -> Result<(), MemoryError> {
        let available = self.total_bytes.saturating_sub(self.allocated());
        if bytes > available {
            return Err(MemoryError::OutOfMemory {
                requested: bytes,
                available,
                total: self.total_bytes,
                allocated: self.allocated(),
            });
        }
        self.weight_bytes += bytes;
        tracing::debug!(
            backend = %self.backend_name,
            weight_bytes = self.weight_bytes,
            total_allocated = self.allocated(),
            "Allocated weight memory"
        );
        Ok(())
    }

    /// Reserve bytes for KV cache. Returns error if it would exceed budget.
    pub fn allocate_kv_cache(&mut self, bytes: usize) -> Result<(), MemoryError> {
        let available = self.total_bytes.saturating_sub(self.allocated());
        if bytes > available {
            return Err(MemoryError::OutOfMemory {
                requested: bytes,
                available,
                total: self.total_bytes,
                allocated: self.allocated(),
            });
        }
        self.kv_cache_bytes += bytes;
        tracing::debug!(
            backend = %self.backend_name,
            kv_cache_bytes = self.kv_cache_bytes,
            total_allocated = self.allocated(),
            "Allocated KV cache memory"
        );
        Ok(())
    }

    /// Release weight memory (e.g. when unloading a model).
    pub fn free_weights(&mut self, bytes: usize) {
        self.weight_bytes = self.weight_bytes.saturating_sub(bytes);
    }

    /// Release KV cache memory (e.g. on sequence reset).
    pub fn free_kv_cache(&mut self, bytes: usize) {
        self.kv_cache_bytes = self.kv_cache_bytes.saturating_sub(bytes);
    }

    /// Release all tracked memory.
    pub fn free_all(&mut self) {
        self.weight_bytes = 0;
        self.kv_cache_bytes = 0;
    }

    /// Current memory usage snapshot.
    pub fn usage(&self) -> MemoryUsage {
        MemoryUsage {
            allocated: self.allocated(),
            total: self.total_bytes,
            free: self.total_bytes.saturating_sub(self.allocated()),
            weights: self.weight_bytes,
            kv_cache: self.kv_cache_bytes,
        }
    }

    /// Check whether `bytes` can be allocated without exceeding budget.
    pub fn can_allocate(&self, bytes: usize) -> bool {
        bytes <= self.total_bytes.saturating_sub(self.allocated())
    }

    /// Estimate weight memory for a model config (all weights dequantized to f32).
    ///
    /// This is an estimate based on model dimensions. Actual memory may differ
    /// due to alignment, padding, or tied embeddings.
    pub fn estimate_weight_bytes(config: &crate::model::config::LLMHostsModelConfig) -> usize {
        let f32_size = std::mem::size_of::<f32>();
        let h = config.hidden_size;
        let kv_dim = config.num_kv_heads * config.head_dim;
        let ffn = config.intermediate_size;
        let v = config.vocab_size;

        // token_embd: [vocab, hidden]
        let embd = v * h * f32_size;
        // output_norm: [hidden]
        let out_norm = h * f32_size;
        // output (LM head): [hidden, vocab] (transposed at load time)
        let output = h * v * f32_size;

        // Per-layer weights:
        // attn_norm: [hidden]
        // attn_q: [hidden, hidden] (transposed)
        // attn_k: [hidden, kv_dim] (transposed)
        // attn_v: [hidden, kv_dim] (transposed)
        // attn_output: [hidden, hidden] (transposed)
        // ffn_norm: [hidden]
        // ffn_gate: [hidden, ffn] (transposed)
        // ffn_up: [hidden, ffn] (transposed)
        // ffn_down: [ffn, hidden] (transposed)
        let per_layer = (h          // attn_norm
            + h * h                 // attn_q
            + h * kv_dim            // attn_k
            + h * kv_dim            // attn_v
            + h * h                 // attn_output
            + h                     // ffn_norm
            + h * ffn               // ffn_gate
            + h * ffn               // ffn_up
            + ffn * h)              // ffn_down
            * f32_size;

        embd + out_norm + output + per_layer * config.num_layers
    }

    /// Estimate KV cache memory for a model config at max sequence length.
    ///
    /// Each layer has a K and V buffer of [max_seq_len * kv_dim] f32 values.
    pub fn estimate_kv_cache_bytes(config: &crate::model::config::LLMHostsModelConfig) -> usize {
        let f32_size = std::mem::size_of::<f32>();
        let kv_dim = config.num_kv_heads * config.head_dim;
        // Per layer: K buffer + V buffer
        let per_layer = 2 * config.max_seq_len * kv_dim * f32_size;
        per_layer * config.num_layers
    }
}

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

    #[test]
    fn test_new_allocator_starts_empty() {
        let alloc = GpuAllocator::new("cuda:0", 1_000_000);
        assert_eq!(alloc.backend_name(), "cuda:0");
        assert_eq!(alloc.total_bytes(), 1_000_000);
        let usage = alloc.usage();
        assert_eq!(usage.allocated, 0);
        assert_eq!(usage.free, 1_000_000);
        assert_eq!(usage.weights, 0);
        assert_eq!(usage.kv_cache, 0);
    }

    #[test]
    fn test_allocate_weights_tracks_correctly() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(400).unwrap();
        let usage = alloc.usage();
        assert_eq!(usage.weights, 400);
        assert_eq!(usage.allocated, 400);
        assert_eq!(usage.free, 600);
    }

    #[test]
    fn test_allocate_kv_cache_tracks_correctly() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_kv_cache(300).unwrap();
        let usage = alloc.usage();
        assert_eq!(usage.kv_cache, 300);
        assert_eq!(usage.allocated, 300);
        assert_eq!(usage.free, 700);
    }

    #[test]
    fn test_combined_allocation() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(400).unwrap();
        alloc.allocate_kv_cache(300).unwrap();
        let usage = alloc.usage();
        assert_eq!(usage.weights, 400);
        assert_eq!(usage.kv_cache, 300);
        assert_eq!(usage.allocated, 700);
        assert_eq!(usage.free, 300);
    }

    #[test]
    fn test_oom_on_weight_allocation() {
        let mut alloc = GpuAllocator::new("cuda:0", 500);
        alloc.allocate_weights(300).unwrap();
        let err = alloc.allocate_weights(300).unwrap_err();
        assert_eq!(
            err,
            MemoryError::OutOfMemory {
                requested: 300,
                available: 200,
                total: 500,
                allocated: 300,
            }
        );
        // Allocation unchanged after failure
        assert_eq!(alloc.usage().weights, 300);
    }

    #[test]
    fn test_oom_on_kv_cache_allocation() {
        let mut alloc = GpuAllocator::new("cuda:0", 500);
        alloc.allocate_weights(400).unwrap();
        let err = alloc.allocate_kv_cache(200).unwrap_err();
        assert_eq!(
            err,
            MemoryError::OutOfMemory {
                requested: 200,
                available: 100,
                total: 500,
                allocated: 400,
            }
        );
        assert_eq!(alloc.usage().kv_cache, 0);
    }

    #[test]
    fn test_free_weights() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(400).unwrap();
        alloc.free_weights(200);
        assert_eq!(alloc.usage().weights, 200);
        assert_eq!(alloc.usage().free, 800);
    }

    #[test]
    fn test_free_kv_cache() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_kv_cache(400).unwrap();
        alloc.free_kv_cache(400);
        assert_eq!(alloc.usage().kv_cache, 0);
        assert_eq!(alloc.usage().free, 1_000);
    }

    #[test]
    fn test_free_saturates_at_zero() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(100).unwrap();
        alloc.free_weights(500); // more than allocated
        assert_eq!(alloc.usage().weights, 0);
    }

    #[test]
    fn test_free_all() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(300).unwrap();
        alloc.allocate_kv_cache(200).unwrap();
        alloc.free_all();
        let usage = alloc.usage();
        assert_eq!(usage.allocated, 0);
        assert_eq!(usage.weights, 0);
        assert_eq!(usage.kv_cache, 0);
        assert_eq!(usage.free, 1_000);
    }

    #[test]
    fn test_can_allocate() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        assert!(alloc.can_allocate(1_000));
        assert!(!alloc.can_allocate(1_001));
        alloc.allocate_weights(600).unwrap();
        assert!(alloc.can_allocate(400));
        assert!(!alloc.can_allocate(401));
    }

    #[test]
    fn test_exact_budget_allocation() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(500).unwrap();
        alloc.allocate_kv_cache(500).unwrap();
        assert_eq!(alloc.usage().free, 0);
        assert!(!alloc.can_allocate(1));
    }

    #[test]
    fn test_zero_budget_allocator() {
        let mut alloc = GpuAllocator::new("cpu", 0);
        let err = alloc.allocate_weights(1).unwrap_err();
        match err {
            MemoryError::OutOfMemory { requested, available, .. } => {
                assert_eq!(requested, 1);
                assert_eq!(available, 0);
            }
        }
    }

    #[test]
    fn test_estimate_weight_bytes() {
        let config = crate::model::config::LLMHostsModelConfig {
            architecture: "llama".to_string(),
            hidden_size: 64,
            num_attention_heads: 4,
            num_kv_heads: 2,
            num_layers: 2,
            intermediate_size: 128,
            vocab_size: 128,
            max_seq_len: 256,
            head_dim: 16,
            rope_freq_base: 10000.0,
            rms_norm_eps: 1e-5,
            eos_token_id: None,
            bos_token_id: None,
        };
        let bytes = GpuAllocator::estimate_weight_bytes(&config);
        // Should be positive and reasonable
        assert!(bytes > 0);
        // token_embd: 128*64*4 = 32768
        // output_norm: 64*4 = 256
        // output: 64*128*4 = 32768
        // Per layer (kv_dim = 2*16 = 32):
        //   attn_norm: 64 * 4 = 256
        //   attn_q: 64*64 * 4 = 16384
        //   attn_k: 64*32 * 4 = 8192
        //   attn_v: 64*32 * 4 = 8192
        //   attn_output: 64*64 * 4 = 16384
        //   ffn_norm: 64 * 4 = 256
        //   ffn_gate: 64*128 * 4 = 32768
        //   ffn_up: 64*128 * 4 = 32768
        //   ffn_down: 128*64 * 4 = 32768
        // per_layer = (64 + 4096 + 2048 + 2048 + 4096 + 64 + 8192 + 8192 + 8192) * 4 = 147968
        // total = 32768 + 256 + 32768 + 147968 * 2 = 361728
        assert_eq!(bytes, 361_728);
    }

    #[test]
    fn test_estimate_kv_cache_bytes() {
        let config = crate::model::config::LLMHostsModelConfig {
            architecture: "llama".to_string(),
            hidden_size: 64,
            num_attention_heads: 4,
            num_kv_heads: 2,
            num_layers: 2,
            intermediate_size: 128,
            vocab_size: 128,
            max_seq_len: 256,
            head_dim: 16,
            rope_freq_base: 10000.0,
            rms_norm_eps: 1e-5,
            eos_token_id: None,
            bos_token_id: None,
        };
        let bytes = GpuAllocator::estimate_kv_cache_bytes(&config);
        // kv_dim = 2 * 16 = 32
        // per_layer = 2 * 256 * 32 * 4 = 65536
        // total = 65536 * 2 = 131072
        assert_eq!(bytes, 131_072);
    }

    #[test]
    fn test_memory_usage_display() {
        let usage = MemoryUsage {
            allocated: 1_048_576, // 1 MB
            total: 4_194_304,     // 4 MB
            free: 3_145_728,      // 3 MB
            weights: 786_432,     // 0.75 MB
            kv_cache: 262_144,    // 0.25 MB
        };
        let s = format!("{usage}");
        assert!(s.contains("1.0"));
        assert!(s.contains("4.0"));
        assert!(s.contains("free"));
    }

    #[test]
    fn test_allocate_then_free_then_reuse() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000);
        alloc.allocate_weights(800).unwrap();
        assert!(!alloc.can_allocate(300));
        alloc.free_weights(800);
        assert!(alloc.can_allocate(300));
        alloc.allocate_kv_cache(1_000).unwrap();
        assert_eq!(alloc.usage().free, 0);
    }

    #[test]
    fn test_error_message_is_descriptive() {
        let mut alloc = GpuAllocator::new("cuda:0", 1_000_000);
        alloc.allocate_weights(900_000).unwrap();
        let err = alloc.allocate_kv_cache(200_000).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("Out of memory"));
        assert!(msg.contains("200000"));
        assert!(msg.contains("100000")); // available
    }
}
