//! Optimized Buffer Pool with Size Classes
//!
//! Provides efficient buffer allocation using size-class bucketing,
//! reducing fragmentation and improving reuse rates.
//!
//! # Features
//! - Size-class bucketing (power-of-2 sizes)
//! - Per-bucket statistics
//! - Configurable retention policy
//! - Memory pressure handling

use crate::vulkan::Allocator;
use ash::vk;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

/// Size class configuration
const MIN_SIZE_CLASS: u32 = 8; // 256 bytes (2^8)
const MAX_SIZE_CLASS: u32 = 26; // 64 MB (2^26)
const NUM_SIZE_CLASSES: usize = (MAX_SIZE_CLASS - MIN_SIZE_CLASS + 1) as usize;

/// Get size class index for a given size
fn size_class_index(size: u64) -> usize {
    if size == 0 {
        return 0;
    }
    // Find the power of 2 that fits this size
    let bits = 64 - size.saturating_sub(1).leading_zeros();
    let class = bits.clamp(MIN_SIZE_CLASS, MAX_SIZE_CLASS);
    (class - MIN_SIZE_CLASS) as usize
}

/// Get actual allocation size for a size class
fn size_class_size(index: usize) -> u64 {
    1u64 << (index as u32 + MIN_SIZE_CLASS)
}

/// Buffer allocation with additional metadata
#[derive(Debug, Clone)]
pub struct BufferAllocation {
    pub buffer: vk::Buffer,
    pub size: u64,
    pub actual_size: u64, // Actual allocated size (may be larger)
    pub offset: u64,
    pub name: Option<String>,
    pub size_class: usize,
    pub frame_last_used: u64,
}

/// Per-size-class statistics
#[derive(Debug, Clone, Default)]
pub struct SizeClassStats {
    pub allocations: u64,
    pub deallocations: u64,
    pub reuses: u64,
    pub current_available: usize,
    pub current_in_use: usize,
    pub total_bytes: u64,
}

/// Pool-wide statistics
#[derive(Debug, Clone, Default)]
pub struct BufferPoolStats {
    pub total_allocations: u64,
    pub total_deallocations: u64,
    pub total_reuses: u64,
    pub current_available: usize,
    pub current_in_use: usize,
    pub total_allocated_bytes: u64,
    pub reuse_rate: f64,
    pub size_classes: Vec<SizeClassStats>,
}

impl BufferPoolStats {
    /// Format stats as a summary string
    pub fn format(&self) -> String {
        format!(
            "BufferPool: {} allocs ({:.1}% reuse), {} available, {} in use, {:.2} MB",
            self.total_allocations,
            self.reuse_rate * 100.0,
            self.current_available,
            self.current_in_use,
            self.total_allocated_bytes as f64 / (1024.0 * 1024.0)
        )
    }
}

/// Per-size-class bucket
struct SizeClassBucket {
    available: VecDeque<BufferAllocation>,
    in_use: Vec<BufferAllocation>,
    stats: SizeClassStats,
}

impl SizeClassBucket {
    fn new() -> Self {
        Self {
            available: VecDeque::new(),
            in_use: Vec::new(),
            stats: SizeClassStats::default(),
        }
    }
}

/// Configuration for the buffer pool
#[derive(Debug, Clone)]
pub struct BufferPoolConfig {
    /// Maximum buffers to keep in each size class
    pub max_per_class: usize,
    /// Frames before releasing unused buffers
    pub retention_frames: u64,
    /// Enable aggressive memory reclamation
    pub aggressive_reclaim: bool,
}

impl Default for BufferPoolConfig {
    fn default() -> Self {
        Self {
            max_per_class: 16,
            retention_frames: 300, // ~5 seconds at 60fps
            aggressive_reclaim: false,
        }
    }
}

/// Optimized buffer pool with size-class bucketing
pub struct BufferPool {
    allocator: Arc<Allocator>,
    buckets: Mutex<Vec<SizeClassBucket>>,
    config: BufferPoolConfig,
    current_frame: AtomicU64,
    // Global stats (atomic for lock-free reads)
    total_allocations: AtomicU64,
    total_reuses: AtomicU64,
    total_allocated_bytes: AtomicU64,
}

impl BufferPool {
    /// Creates a new optimized buffer pool
    pub fn new(allocator: Arc<Allocator>) -> Self {
        Self::with_config(allocator, BufferPoolConfig::default())
    }

    /// Creates a buffer pool with custom configuration
    pub fn with_config(allocator: Arc<Allocator>, config: BufferPoolConfig) -> Self {
        let buckets = (0..NUM_SIZE_CLASSES)
            .map(|_| SizeClassBucket::new())
            .collect();

        Self {
            allocator,
            buckets: Mutex::new(buckets),
            config,
            current_frame: AtomicU64::new(0),
            total_allocations: AtomicU64::new(0),
            total_reuses: AtomicU64::new(0),
            total_allocated_bytes: AtomicU64::new(0),
        }
    }

    /// Advance to next frame (call once per frame)
    pub fn next_frame(&self) {
        self.current_frame.fetch_add(1, Ordering::Relaxed);
    }

    /// Get current frame number
    pub fn current_frame(&self) -> u64 {
        self.current_frame.load(Ordering::Relaxed)
    }

    /// Allocates a buffer from the appropriate size class
    ///
    /// # Safety
    /// Same as original allocate - caller must ensure valid usage flags
    pub unsafe fn allocate(
        &self,
        size: u64,
        usage: vk::BufferUsageFlags,
        memory_usage: vk_mem::MemoryUsage,
        name: Option<String>,
    ) -> crate::Result<BufferAllocation> {
        let class_index = size_class_index(size);
        let actual_size = size_class_size(class_index);
        let frame = self.current_frame.load(Ordering::Relaxed);

        let mut buckets = self.buckets.lock().unwrap();
        let bucket = &mut buckets[class_index];

        // Try to find a reusable buffer in this size class
        if let Some(mut alloc) = bucket.available.pop_front() {
            alloc.size = size;
            alloc.name = name.clone();
            alloc.frame_last_used = frame;
            bucket.in_use.push(alloc.clone());
            bucket.stats.reuses += 1;
            self.total_reuses.fetch_add(1, Ordering::Relaxed);

            if let Some(ref n) = name {
                log::trace!("Reusing buffer '{n}' (class {class_index}, {actual_size} bytes)");
            }
            return Ok(alloc);
        }

        // Need to allocate new buffer
        drop(buckets); // Release lock during expensive allocation

        if let Some(ref n) = name {
            log::debug!("Allocating new buffer '{n}' (class {class_index}, {actual_size} bytes)");
        }

        let (buffer, _allocation) =
            self.allocator
                .create_buffer(actual_size, usage, memory_usage)?;

        self.total_allocations.fetch_add(1, Ordering::Relaxed);
        self.total_allocated_bytes
            .fetch_add(actual_size, Ordering::Relaxed);

        let alloc = BufferAllocation {
            buffer,
            size,
            actual_size,
            offset: 0,
            name,
            size_class: class_index,
            frame_last_used: frame,
        };

        // Re-acquire lock to update tracking
        let mut buckets = self.buckets.lock().unwrap();
        buckets[class_index].in_use.push(alloc.clone());
        buckets[class_index].stats.allocations += 1;
        buckets[class_index].stats.total_bytes += actual_size;

        Ok(alloc)
    }

    /// Returns a buffer to the pool
    pub fn deallocate(&self, buffer: BufferAllocation) {
        let class_index = buffer.size_class;
        let mut buckets = self.buckets.lock().unwrap();
        let bucket = &mut buckets[class_index];

        // Remove from in_use
        bucket.in_use.retain(|b| b.buffer != buffer.buffer);

        // Validate retention criteria.
        if bucket.available.len() < self.config.max_per_class {
            if let Some(ref name) = buffer.name {
                log::trace!("Returning buffer '{name}' to pool (class {class_index})");
            }
            bucket.available.push_back(buffer);
        } else {
            // Pool is full for this class, actually free the buffer
            // Buffers are destroyed via the allocator instance.
            log::trace!("Pool full for class {class_index}, buffer will be dropped");
        }

        bucket.stats.deallocations += 1;
    }

    /// Reclaim memory from unused buffers
    pub fn reclaim_memory(&self) {
        let frame = self.current_frame.load(Ordering::Relaxed);
        let retention = self.config.retention_frames;

        let mut buckets = self.buckets.lock().unwrap();
        for bucket in buckets.iter_mut() {
            bucket
                .available
                .retain(|alloc| frame.saturating_sub(alloc.frame_last_used) < retention);
        }
    }

    /// Get comprehensive statistics
    pub fn stats(&self) -> BufferPoolStats {
        let buckets = self.buckets.lock().unwrap();

        let mut stats = BufferPoolStats {
            total_allocations: self.total_allocations.load(Ordering::Relaxed),
            total_deallocations: 0,
            total_reuses: self.total_reuses.load(Ordering::Relaxed),
            current_available: 0,
            current_in_use: 0,
            total_allocated_bytes: self.total_allocated_bytes.load(Ordering::Relaxed),
            reuse_rate: 0.0,
            size_classes: Vec::with_capacity(NUM_SIZE_CLASSES),
        };

        for bucket in buckets.iter() {
            stats.current_available += bucket.available.len();
            stats.current_in_use += bucket.in_use.len();
            stats.total_deallocations += bucket.stats.deallocations;

            let mut class_stats = bucket.stats.clone();
            class_stats.current_available = bucket.available.len();
            class_stats.current_in_use = bucket.in_use.len();
            stats.size_classes.push(class_stats);
        }

        let total_requests = stats.total_allocations + stats.total_reuses;
        if total_requests > 0 {
            stats.reuse_rate = stats.total_reuses as f64 / total_requests as f64;
        }

        stats
    }

    /// Get simple stats tuple (legacy API compatibility)
    pub fn simple_stats(&self) -> (usize, usize, u64) {
        let buckets = self.buckets.lock().unwrap();
        let mut available = 0;
        let mut in_use = 0;
        for bucket in buckets.iter() {
            available += bucket.available.len();
            in_use += bucket.in_use.len();
        }
        (
            available,
            in_use,
            self.total_allocated_bytes.load(Ordering::Relaxed),
        )
    }
}

impl Drop for BufferPool {
    fn drop(&mut self) {
        let stats = self.stats();
        log::info!(
            "Buffer pool destroyed: {} reuses/{} allocs ({:.1}% reuse), {:.2} MB allocated",
            stats.total_reuses,
            stats.total_allocations,
            stats.reuse_rate * 100.0,
            stats.total_allocated_bytes as f64 / (1024.0 * 1024.0)
        );
    }
}

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

    #[test]
    fn test_size_class_index() {
        // 256 bytes = class 0
        assert_eq!(size_class_index(256), 0);
        // 257 bytes = class 1 (512 bytes)
        assert_eq!(size_class_index(257), 1);
        // 1KB = class 2
        assert_eq!(size_class_index(1024), 2);
        // 1MB = class 12
        assert_eq!(size_class_index(1024 * 1024), 12);
    }

    #[test]
    fn test_size_class_size() {
        assert_eq!(size_class_size(0), 256);
        assert_eq!(size_class_size(2), 1024);
        assert_eq!(size_class_size(12), 1024 * 1024);
    }

    #[test]
    fn test_stats_format() {
        let stats = BufferPoolStats {
            total_allocations: 100,
            total_reuses: 80,
            reuse_rate: 0.8,
            current_available: 10,
            current_in_use: 5,
            total_allocated_bytes: 1024 * 1024,
            ..Default::default()
        };
        let formatted = stats.format();
        assert!(formatted.contains("80.0% reuse"));
        assert!(formatted.contains("1.00 MB"));
    }
}
