//! GPU memory pool and OOM handling
//!
//! Provides a slab-based GPU memory allocator with automatic OOM detection
//! and graceful degradation to CPU when GPU memory is exhausted.
//!
//! ## Design Goals
//!
//! 1. **Pre-allocation**: Reduce CUDA malloc overhead during index building
//! 2. **OOM Detection**: Detect out-of-memory before it causes a crash
//! 3. **Auto-degradation**: Spill to host pinned memory when GPU memory is full
//! 4. **Accounting**: Track peak memory usage for benchmarking

use std::sync::atomic::{AtomicUsize, Ordering};
use super::error::{GpuError, GpuResult};
use super::device::GpuDeviceInfo;

/// Default slab size for GPU memory pool (2 MB).
pub const DEFAULT_SLAB_SIZE: usize = 2 * 1024 * 1024;

/// Default memory budget fraction (use 80% of free GPU memory).
pub const DEFAULT_BUDGET_FRACTION: f64 = 0.80;

// ============================================================================
// GpuMemoryPool — slab-based GPU memory manager
// ============================================================================

/// GPU memory pool with OOM detection and auto-degradation.
///
/// The pool pre-allocates a budget of GPU memory and tracks usage.
/// When the budget is exceeded, it can:
/// 1. Spill to host (CPU) pinned memory
/// 2. Signal the caller to reduce batch size
/// 3. Trigger a full fallback to CPU pipeline
///
/// ## Thread Safety
///
/// Memory accounting is atomic, but actual allocations should be
/// done from a single thread (the GPU worker thread).
pub struct GpuMemoryPool {
    /// Device this pool is associated with
    device_info: GpuDeviceInfo,

    /// Total budget in bytes
    budget_bytes: usize,

    /// Currently allocated bytes (atomic for concurrent reads)
    allocated_bytes: AtomicUsize,

    /// Peak allocated bytes (for reporting)
    peak_bytes: AtomicUsize,

    /// Number of allocations performed
    allocation_count: AtomicUsize,

    /// Number of times OOM was avoided by degradation
    oom_avoidance_count: AtomicUsize,

    /// Slab size for allocations
    slab_size: usize,
}

/// Statistics about GPU memory pool usage.
#[derive(Debug, Clone)]
pub struct GpuMemoryStats {
    /// Device name
    pub device_name: String,
    /// Total budget in bytes
    pub budget_bytes: usize,
    /// Currently allocated bytes
    pub allocated_bytes: usize,
    /// Peak allocated bytes
    pub peak_bytes: usize,
    /// Number of allocations
    pub allocation_count: usize,
    /// Number of OOM avoidances
    pub oom_avoidance_count: usize,
    /// Budget utilization (0.0 - 1.0)
    pub utilization: f64,
}

/// Result of a memory reservation request.
#[derive(Debug)]
pub enum MemoryReservation {
    /// Memory reserved on GPU successfully.
    Gpu {
        /// Size in bytes that was reserved
        size_bytes: usize,
    },

    /// GPU memory insufficient, spilled to CPU host memory.
    /// The caller should use CPU buffers instead of GPU buffers.
    SpilledToCpu {
        /// Requested size in bytes
        requested_bytes: usize,
        /// Available GPU memory at time of request
        available_gpu_bytes: usize,
        /// Reason for spill
        reason: String,
    },
}

impl MemoryReservation {
    /// Whether this reservation is on GPU.
    pub fn is_gpu(&self) -> bool {
        matches!(self, MemoryReservation::Gpu { .. })
    }

    /// Whether this reservation spilled to CPU.
    pub fn is_cpu_spill(&self) -> bool {
        matches!(self, MemoryReservation::SpilledToCpu { .. })
    }
}

impl GpuMemoryPool {
    /// Create a new GPU memory pool with auto-calculated budget.
    ///
    /// Budget = min(explicit_budget, device_free_memory * DEFAULT_BUDGET_FRACTION)
    pub fn new(device_info: GpuDeviceInfo, explicit_budget: Option<usize>) -> Self {
        let auto_budget =
            (device_info.free_memory_bytes as f64 * DEFAULT_BUDGET_FRACTION) as usize;
        let budget = explicit_budget
            .map(|b| b.min(auto_budget))
            .unwrap_or(auto_budget);

        tracing::info!(
            "GPU memory pool: budget={:.1}MB (device {} has {:.1}MB free)",
            budget as f64 / (1024.0 * 1024.0),
            device_info.name,
            device_info.free_memory_bytes as f64 / (1024.0 * 1024.0),
        );

        Self {
            device_info,
            budget_bytes: budget,
            allocated_bytes: AtomicUsize::new(0),
            peak_bytes: AtomicUsize::new(0),
            allocation_count: AtomicUsize::new(0),
            oom_avoidance_count: AtomicUsize::new(0),
            slab_size: DEFAULT_SLAB_SIZE,
        }
    }

    /// Create a pool with a specific slab size.
    pub fn with_slab_size(mut self, slab_size: usize) -> Self {
        self.slab_size = slab_size;
        self
    }

    /// Try to reserve GPU memory for a workload.
    ///
    /// If there's enough budget, returns `Gpu` reservation.
    /// If budget would be exceeded, returns `SpilledToCpu` (caller should use CPU buffers).
    ///
    /// This does NOT actually allocate CUDA memory — it's a budget check.
    /// Actual CUDA allocation happens when the caller creates tensors/buffers.
    pub fn try_reserve(&self, size_bytes: usize) -> MemoryReservation {
        let current = self.allocated_bytes.load(Ordering::Relaxed);
        let new_total = current.saturating_add(size_bytes);

        if new_total <= self.budget_bytes {
            // Reserve on GPU
            self.allocated_bytes.store(new_total, Ordering::Relaxed);
            self.allocation_count.fetch_add(1, Ordering::Relaxed);

            // Update peak
            let mut peak = self.peak_bytes.load(Ordering::Relaxed);
            while new_total > peak {
                match self.peak_bytes.compare_exchange_weak(
                    peak,
                    new_total,
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => break,
                    Err(actual) => peak = actual,
                }
            }

            MemoryReservation::Gpu { size_bytes }
        } else {
            // Spill to CPU
            self.oom_avoidance_count.fetch_add(1, Ordering::Relaxed);
            let available = self.budget_bytes.saturating_sub(current);

            tracing::warn!(
                "GPU memory budget exceeded: requested {:.1}MB, available {:.1}MB — spilling to CPU",
                size_bytes as f64 / (1024.0 * 1024.0),
                available as f64 / (1024.0 * 1024.0),
            );

            MemoryReservation::SpilledToCpu {
                requested_bytes: size_bytes,
                available_gpu_bytes: available,
                reason: format!(
                    "Budget {:.1}MB would be exceeded (current {:.1}MB + request {:.1}MB)",
                    self.budget_bytes as f64 / (1024.0 * 1024.0),
                    current as f64 / (1024.0 * 1024.0),
                    size_bytes as f64 / (1024.0 * 1024.0),
                ),
            }
        }
    }

    /// Release previously reserved GPU memory.
    pub fn release(&self, size_bytes: usize) {
        let prev = self.allocated_bytes.fetch_sub(
            size_bytes.min(self.allocated_bytes.load(Ordering::Relaxed)),
            Ordering::Relaxed,
        );
        tracing::trace!(
            "GPU memory released: {:.1}MB (was {:.1}MB, now {:.1}MB)",
            size_bytes as f64 / (1024.0 * 1024.0),
            prev as f64 / (1024.0 * 1024.0),
            self.allocated_bytes.load(Ordering::Relaxed) as f64 / (1024.0 * 1024.0),
        );
    }

    /// Check if a given size would fit in the remaining budget.
    pub fn can_fit(&self, size_bytes: usize) -> bool {
        let current = self.allocated_bytes.load(Ordering::Relaxed);
        current.saturating_add(size_bytes) <= self.budget_bytes
    }

    /// Calculate optimal batch size for a given vector dimension.
    ///
    /// Returns the number of vectors per batch that fits in 50% of remaining budget
    /// (leaving room for intermediate computations).
    pub fn optimal_batch_size(&self, dimension: usize) -> usize {
        let current = self.allocated_bytes.load(Ordering::Relaxed);
        let remaining = self.budget_bytes.saturating_sub(current);
        // Use 50% of remaining for the batch (rest for intermediate results)
        let batch_budget = remaining / 2;
        let bytes_per_vector = dimension * std::mem::size_of::<f32>();
        if bytes_per_vector == 0 {
            return 0;
        }
        // Round down to slab_size granularity
        let raw_count = batch_budget / bytes_per_vector;
        // Minimum batch of 1024, maximum of 1M
        raw_count.max(1024).min(1_000_000)
    }

    /// Get current memory statistics.
    pub fn stats(&self) -> GpuMemoryStats {
        let allocated = self.allocated_bytes.load(Ordering::Relaxed);
        GpuMemoryStats {
            device_name: self.device_info.name.clone(),
            budget_bytes: self.budget_bytes,
            allocated_bytes: allocated,
            peak_bytes: self.peak_bytes.load(Ordering::Relaxed),
            allocation_count: self.allocation_count.load(Ordering::Relaxed),
            oom_avoidance_count: self.oom_avoidance_count.load(Ordering::Relaxed),
            utilization: if self.budget_bytes > 0 {
                allocated as f64 / self.budget_bytes as f64
            } else {
                0.0
            },
        }
    }

    /// Get the total budget in bytes.
    pub fn budget(&self) -> usize {
        self.budget_bytes
    }

    /// Get currently allocated bytes.
    pub fn allocated(&self) -> usize {
        self.allocated_bytes.load(Ordering::Relaxed)
    }

    /// Get peak allocated bytes.
    pub fn peak(&self) -> usize {
        self.peak_bytes.load(Ordering::Relaxed)
    }

    /// Reset allocation counters (for benchmarking).
    pub fn reset_stats(&self) {
        self.allocated_bytes.store(0, Ordering::Relaxed);
        self.peak_bytes.store(0, Ordering::Relaxed);
        self.allocation_count.store(0, Ordering::Relaxed);
        self.oom_avoidance_count.store(0, Ordering::Relaxed);
    }

    /// Check memory health and return any warnings.
    pub fn health_check(&self) -> GpuResult<()> {
        let stats = self.stats();

        if stats.utilization > 0.95 {
            return Err(GpuError::OutOfMemory {
                requested_bytes: 0,
                available_bytes: self.budget_bytes.saturating_sub(stats.allocated_bytes),
                device_name: self.device_info.name.clone(),
            });
        }

        Ok(())
    }
}

impl std::fmt::Display for GpuMemoryStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "GPU Memory [{}]: {:.1}MB / {:.1}MB ({:.1}% used), peak {:.1}MB, {} allocs, {} OOM avoided",
            self.device_name,
            self.allocated_bytes as f64 / (1024.0 * 1024.0),
            self.budget_bytes as f64 / (1024.0 * 1024.0),
            self.utilization * 100.0,
            self.peak_bytes as f64 / (1024.0 * 1024.0),
            self.allocation_count,
            self.oom_avoidance_count,
        )
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    fn test_device_info() -> GpuDeviceInfo {
        GpuDeviceInfo {
            device_id: 0,
            name: "Test GPU".into(),
            compute_capability: (8, 0),
            total_memory_bytes: 8 * 1024 * 1024 * 1024,    // 8GB
            free_memory_bytes: 6 * 1024 * 1024 * 1024,     // 6GB free
            sm_count: 80,
            is_compatible: true,
        }
    }

    #[test]
    fn test_pool_creation() {
        let pool = GpuMemoryPool::new(test_device_info(), None);
        // Budget should be 80% of 6GB = 4.8GB
        let expected = (6.0 * 1024.0 * 1024.0 * 1024.0 * 0.8) as usize;
        assert_eq!(pool.budget(), expected);
        assert_eq!(pool.allocated(), 0);
        assert_eq!(pool.peak(), 0);
    }

    #[test]
    fn test_pool_explicit_budget() {
        let budget = 1024 * 1024 * 1024; // 1GB
        let pool = GpuMemoryPool::new(test_device_info(), Some(budget));
        assert_eq!(pool.budget(), budget);
    }

    #[test]
    fn test_pool_explicit_budget_capped() {
        // If explicit budget > auto budget, auto budget wins
        let huge_budget = 100 * 1024 * 1024 * 1024; // 100GB (exceeds device)
        let pool = GpuMemoryPool::new(test_device_info(), Some(huge_budget));
        let expected = (6.0 * 1024.0 * 1024.0 * 1024.0 * 0.8) as usize;
        assert_eq!(pool.budget(), expected);
    }

    #[test]
    fn test_reserve_and_release() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(100 * 1024 * 1024)); // 100MB

        // Reserve 30MB
        let res = pool.try_reserve(30 * 1024 * 1024);
        assert!(res.is_gpu());
        assert_eq!(pool.allocated(), 30 * 1024 * 1024);

        // Reserve another 30MB
        let res = pool.try_reserve(30 * 1024 * 1024);
        assert!(res.is_gpu());
        assert_eq!(pool.allocated(), 60 * 1024 * 1024);

        // Release 30MB
        pool.release(30 * 1024 * 1024);
        assert_eq!(pool.allocated(), 30 * 1024 * 1024);
        assert_eq!(pool.peak(), 60 * 1024 * 1024);
    }

    #[test]
    fn test_oom_spill() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(100 * 1024 * 1024)); // 100MB

        // Reserve 80MB
        let res = pool.try_reserve(80 * 1024 * 1024);
        assert!(res.is_gpu());

        // Try to reserve 30MB more (would exceed 100MB budget)
        let res = pool.try_reserve(30 * 1024 * 1024);
        assert!(res.is_cpu_spill());

        let stats = pool.stats();
        assert_eq!(stats.oom_avoidance_count, 1);
    }

    #[test]
    fn test_can_fit() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(100 * 1024 * 1024));
        assert!(pool.can_fit(50 * 1024 * 1024));
        assert!(pool.can_fit(100 * 1024 * 1024));
        assert!(!pool.can_fit(101 * 1024 * 1024));

        pool.try_reserve(80 * 1024 * 1024);
        assert!(pool.can_fit(20 * 1024 * 1024));
        assert!(!pool.can_fit(21 * 1024 * 1024));
    }

    #[test]
    fn test_optimal_batch_size() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(1024 * 1024 * 1024)); // 1GB

        // 768-dim f32 vectors
        let batch = pool.optimal_batch_size(768);
        // Budget/2 = 512MB, each vector = 768*4 = 3072 bytes
        // 512MB / 3072 ≈ 174,762
        assert!(batch > 100_000);
        assert!(batch <= 1_000_000);
    }

    #[test]
    fn test_stats_display() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(100 * 1024 * 1024));
        pool.try_reserve(50 * 1024 * 1024);
        let stats = pool.stats();
        let s = format!("{stats}");
        assert!(s.contains("Test GPU"));
        assert!(s.contains("50.0%") || s.contains("50."));
    }

    #[test]
    fn test_health_check() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(100 * 1024 * 1024));
        assert!(pool.health_check().is_ok());

        // Fill to 96%
        pool.try_reserve(96 * 1024 * 1024);
        assert!(pool.health_check().is_err());
    }

    #[test]
    fn test_reset_stats() {
        let pool = GpuMemoryPool::new(test_device_info(), Some(100 * 1024 * 1024));
        pool.try_reserve(50 * 1024 * 1024);
        assert_eq!(pool.allocated(), 50 * 1024 * 1024);

        pool.reset_stats();
        assert_eq!(pool.allocated(), 0);
        assert_eq!(pool.peak(), 0);
    }
}
