//! GPU Memory Pool
//!
//! This module provides efficient GPU memory management through pooling:
//! - Reduces allocation overhead by reusing memory blocks
//! - Supports multiple allocation strategies (best-fit, first-fit)
//! - Provides memory statistics and fragmentation tracking
//! - Thread-safe for concurrent allocations
//! - Automatic defragmentation and compaction
//!
//! Memory pooling can reduce allocation overhead by 10-100x for
//! frequently allocated tensor sizes during training.

use crate::error::{DistributedError, Result};
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;

// =============================================================================
// Configuration
// =============================================================================

/// Memory allocation strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AllocationStrategy {
    /// Best-fit: Find smallest block that fits (less fragmentation)
    BestFit,
    /// First-fit: Use first block that fits (faster)
    FirstFit,
    /// Buddy allocator: Power-of-2 block sizes
    Buddy,
    /// Slab allocator: Fixed-size blocks for common sizes
    Slab,
}

/// Configuration for memory pool
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryPoolConfig {
    /// Initial pool size in bytes
    pub initial_size_bytes: usize,
    /// Maximum pool size in bytes (0 = unlimited)
    pub max_size_bytes: usize,
    /// Allocation strategy
    pub strategy: AllocationStrategy,
    /// Minimum block size in bytes
    pub min_block_size: usize,
    /// Alignment requirement in bytes
    pub alignment: usize,
    /// Enable automatic defragmentation
    pub auto_defrag: bool,
    /// Fragmentation threshold to trigger defrag (0.0-1.0)
    pub defrag_threshold: f32,
    /// Enable memory usage statistics
    pub enable_stats: bool,
    /// Cache freed blocks for reuse
    pub cache_freed_blocks: bool,
    /// Maximum cached blocks per size class
    pub max_cached_per_size: usize,
}

impl Default for MemoryPoolConfig {
    fn default() -> Self {
        Self {
            initial_size_bytes: 256 * 1024 * 1024, // 256MB
            max_size_bytes: 0,                     // Unlimited
            strategy: AllocationStrategy::BestFit,
            min_block_size: 512,
            alignment: 256, // GPU typically needs 256-byte alignment
            auto_defrag: true,
            defrag_threshold: 0.3, // Defrag when 30% fragmented
            enable_stats: true,
            cache_freed_blocks: true,
            max_cached_per_size: 32,
        }
    }
}

// =============================================================================
// Memory Block
// =============================================================================

/// Unique identifier for a memory block
pub type BlockId = u64;

/// State of a memory block
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BlockState {
    /// Block is free and available
    Free,
    /// Block is allocated and in use
    Allocated,
    /// Block is cached (freed but held for reuse)
    Cached,
}

/// A memory block in the pool
#[derive(Debug, Clone)]
pub struct MemoryBlock {
    /// Unique block ID
    pub id: BlockId,
    /// Offset from pool base address
    pub offset: usize,
    /// Size in bytes
    pub size: usize,
    /// Current state
    pub state: BlockState,
    /// Allocation timestamp (if allocated)
    pub allocated_at: Option<Instant>,
    /// Tag for debugging (e.g., "gradient", "activation")
    pub tag: Option<String>,
    /// Device ID
    pub device_id: i32,
}

impl MemoryBlock {
    /// Create a new free block
    pub fn new(id: BlockId, offset: usize, size: usize, device_id: i32) -> Self {
        Self {
            id,
            offset,
            size,
            state: BlockState::Free,
            allocated_at: None,
            tag: None,
            device_id,
        }
    }

    /// Check if block can satisfy allocation request
    pub fn can_satisfy(&self, size: usize, alignment: usize) -> bool {
        if self.state != BlockState::Free && self.state != BlockState::Cached {
            return false;
        }

        let aligned_offset = (self.offset + alignment - 1) & !(alignment - 1);
        let padding = aligned_offset - self.offset;
        self.size >= size + padding
    }

    /// Get aligned offset within this block
    pub fn aligned_offset(&self, alignment: usize) -> usize {
        (self.offset + alignment - 1) & !(alignment - 1)
    }
}

// =============================================================================
// Size Class Cache
// =============================================================================

/// Cache for commonly-used allocation sizes
pub struct SizeClassCache {
    /// Size class -> list of cached block IDs
    cache: HashMap<usize, VecDeque<BlockId>>,
    /// Maximum blocks per size class
    max_per_size: usize,
    /// Size class granularity
    granularity: usize,
}

impl SizeClassCache {
    /// Create a new size class cache
    pub fn new(max_per_size: usize, granularity: usize) -> Self {
        Self {
            cache: HashMap::new(),
            max_per_size,
            granularity,
        }
    }

    /// Get size class for a given size
    fn size_class(&self, size: usize) -> usize {
        // Round up to nearest granularity
        ((size + self.granularity - 1) / self.granularity) * self.granularity
    }

    /// Try to get a cached block for the given size
    pub fn get(&mut self, size: usize) -> Option<BlockId> {
        let size_class = self.size_class(size);
        self.cache.get_mut(&size_class)?.pop_front()
    }

    /// Cache a freed block
    pub fn put(&mut self, size: usize, block_id: BlockId) -> bool {
        let size_class = self.size_class(size);
        let queue = self.cache.entry(size_class).or_default();

        if queue.len() < self.max_per_size {
            queue.push_back(block_id);
            true
        } else {
            false
        }
    }

    /// Clear all cached blocks
    pub fn clear(&mut self) {
        self.cache.clear();
    }

    /// Get total cached block count
    pub fn cached_count(&self) -> usize {
        self.cache
            .values()
            .map(std::collections::VecDeque::len)
            .sum()
    }
}

// =============================================================================
// Memory Pool Statistics
// =============================================================================

/// Statistics for memory pool usage
#[derive(Debug, Default, Clone)]
pub struct MemoryPoolStats {
    /// Total pool size in bytes
    pub total_size_bytes: usize,
    /// Currently allocated bytes
    pub allocated_bytes: usize,
    /// Peak allocated bytes
    pub peak_allocated_bytes: usize,
    /// Free bytes (including fragmented)
    pub free_bytes: usize,
    /// Number of allocations
    pub allocation_count: usize,
    /// Number of deallocations
    pub deallocation_count: usize,
    /// Number of cache hits
    pub cache_hits: usize,
    /// Number of cache misses
    pub cache_misses: usize,
    /// Fragmentation ratio (0.0 = none, 1.0 = fully fragmented)
    pub fragmentation_ratio: f32,
    /// Number of free blocks
    pub free_block_count: usize,
    /// Average allocation size
    pub avg_allocation_size: usize,
    /// Number of defragmentation runs
    pub defrag_count: usize,
}

impl MemoryPoolStats {
    /// Calculate cache hit ratio
    pub fn cache_hit_ratio(&self) -> f32 {
        let total = self.cache_hits + self.cache_misses;
        if total == 0 {
            0.0
        } else {
            self.cache_hits as f32 / total as f32
        }
    }

    /// Calculate utilization ratio
    pub fn utilization(&self) -> f32 {
        if self.total_size_bytes == 0 {
            0.0
        } else {
            self.allocated_bytes as f32 / self.total_size_bytes as f32
        }
    }
}

// =============================================================================
// Memory Pool
// =============================================================================

/// GPU memory pool for efficient allocation
pub struct MemoryPool {
    /// Configuration
    config: MemoryPoolConfig,
    /// Device ID
    device_id: i32,
    /// All memory blocks
    blocks: RwLock<HashMap<BlockId, MemoryBlock>>,
    /// Free blocks sorted by offset (for coalescing)
    free_blocks: RwLock<BTreeMap<usize, BlockId>>,
    /// Size class cache
    size_cache: Mutex<SizeClassCache>,
    /// Next block ID
    next_block_id: AtomicUsize,
    /// Statistics
    stats: RwLock<MemoryPoolStats>,
    /// Base address (simulated - real impl would use CUDA)
    base_address: usize,
    /// Current pool size
    current_size: AtomicUsize,
}

impl MemoryPool {
    /// Create a new memory pool
    pub fn new(config: MemoryPoolConfig, device_id: i32) -> Result<Self> {
        let initial_size = config.initial_size_bytes;
        let size_cache = SizeClassCache::new(config.max_cached_per_size, config.min_block_size);

        let pool = Self {
            config,
            device_id,
            blocks: RwLock::new(HashMap::new()),
            free_blocks: RwLock::new(BTreeMap::new()),
            size_cache: Mutex::new(size_cache),
            next_block_id: AtomicUsize::new(1),
            stats: RwLock::new(MemoryPoolStats::default()),
            base_address: 0x1000_0000, // Simulated base address
            current_size: AtomicUsize::new(initial_size),
        };

        // Initialize with a single free block
        let initial_block = MemoryBlock::new(0, 0, initial_size, device_id);
        {
            let mut blocks = pool.blocks.write();
            let mut free_blocks = pool.free_blocks.write();
            blocks.insert(0, initial_block);
            free_blocks.insert(0, 0);
        }

        // Initialize stats
        {
            let mut stats = pool.stats.write();
            stats.total_size_bytes = initial_size;
            stats.free_bytes = initial_size;
            stats.free_block_count = 1;
        }

        Ok(pool)
    }

    /// Allocate memory from the pool
    pub fn allocate(&self, size: usize) -> Result<MemoryAllocation> {
        self.allocate_with_tag(size, None)
    }

    /// Allocate memory with a debug tag
    pub fn allocate_with_tag(&self, size: usize, tag: Option<String>) -> Result<MemoryAllocation> {
        let aligned_size = self.align_size(size);

        // Try cache first
        if self.config.cache_freed_blocks {
            let mut cache = self.size_cache.lock();
            if let Some(block_id) = cache.get(aligned_size) {
                let mut blocks = self.blocks.write();
                if let Some(block) = blocks.get_mut(&block_id) {
                    if block.state == BlockState::Cached && block.size >= aligned_size {
                        block.state = BlockState::Allocated;
                        block.allocated_at = Some(Instant::now());
                        block.tag = tag;

                        // Update stats
                        let mut stats = self.stats.write();
                        stats.cache_hits += 1;
                        stats.allocation_count += 1;
                        stats.allocated_bytes += block.size;
                        stats.free_bytes -= block.size;
                        stats.peak_allocated_bytes =
                            stats.peak_allocated_bytes.max(stats.allocated_bytes);

                        return Ok(MemoryAllocation {
                            block_id,
                            ptr: self.base_address + block.offset,
                            size: block.size,
                            device_id: self.device_id,
                        });
                    }
                }
            }

            // Update cache miss stat
            let mut stats = self.stats.write();
            stats.cache_misses += 1;
        }

        // Find a suitable block using the configured strategy
        let block_id = self.find_block(aligned_size)?;

        // Allocate from the block
        self.allocate_from_block(block_id, aligned_size, tag)
    }

    /// Find a suitable block for allocation
    fn find_block(&self, size: usize) -> Result<BlockId> {
        match self.config.strategy {
            AllocationStrategy::BestFit => self.find_best_fit(size),
            AllocationStrategy::FirstFit => self.find_first_fit(size),
            AllocationStrategy::Buddy => self.find_buddy_block(size),
            AllocationStrategy::Slab => self.find_slab_block(size),
        }
    }

    /// Best-fit allocation strategy
    fn find_best_fit(&self, size: usize) -> Result<BlockId> {
        let blocks = self.blocks.read();
        let free_blocks = self.free_blocks.read();

        let mut best_id = None;
        let mut best_size = usize::MAX;

        for &block_id in free_blocks.values() {
            if let Some(block) = blocks.get(&block_id) {
                if block.can_satisfy(size, self.config.alignment) && block.size < best_size {
                    best_id = Some(block_id);
                    best_size = block.size;
                }
            }
        }

        best_id.ok_or_else(|| {
            DistributedError::runtime(format!(
                "No suitable block found for {} bytes allocation",
                size
            ))
        })
    }

    /// First-fit allocation strategy
    fn find_first_fit(&self, size: usize) -> Result<BlockId> {
        let blocks = self.blocks.read();
        let free_blocks = self.free_blocks.read();

        for &block_id in free_blocks.values() {
            if let Some(block) = blocks.get(&block_id) {
                if block.can_satisfy(size, self.config.alignment) {
                    return Ok(block_id);
                }
            }
        }

        Err(DistributedError::runtime(format!(
            "No suitable block found for {} bytes allocation",
            size
        )))
    }

    /// Buddy allocator strategy
    fn find_buddy_block(&self, size: usize) -> Result<BlockId> {
        // Round up to power of 2
        let buddy_size = size.next_power_of_two();
        self.find_best_fit(buddy_size)
    }

    /// Slab allocator strategy
    fn find_slab_block(&self, size: usize) -> Result<BlockId> {
        // Round up to slab size class
        let slab_size = self.align_size(size);
        self.find_best_fit(slab_size)
    }

    /// Allocate from a specific block
    fn allocate_from_block(
        &self,
        block_id: BlockId,
        size: usize,
        tag: Option<String>,
    ) -> Result<MemoryAllocation> {
        let mut blocks = self.blocks.write();
        let mut free_blocks = self.free_blocks.write();

        // First, read the block info we need
        let (aligned_offset, block_offset, block_orig_size) = {
            let block = blocks
                .get(&block_id)
                .ok_or_else(|| DistributedError::runtime("Block not found"))?;

            if block.state != BlockState::Free {
                return Err(DistributedError::runtime("Block is not free"));
            }

            (
                block.aligned_offset(self.config.alignment),
                block.offset,
                block.size,
            )
        };

        let remaining = block_orig_size - size;

        // Split block if there's enough remaining space
        if remaining >= self.config.min_block_size {
            let new_block_id = self.next_block_id.fetch_add(1, Ordering::SeqCst) as BlockId;
            let new_block =
                MemoryBlock::new(new_block_id, block_offset + size, remaining, self.device_id);

            free_blocks.insert(new_block.offset, new_block_id);
            blocks.insert(new_block_id, new_block);
        }

        // Now update the original block
        let alloc_size = if remaining >= self.config.min_block_size {
            size
        } else {
            block_orig_size
        };
        if let Some(block) = blocks.get_mut(&block_id) {
            if remaining >= self.config.min_block_size {
                block.size = size;
            }
            block.state = BlockState::Allocated;
            block.allocated_at = Some(Instant::now());
            block.tag = tag;
        }

        // Remove from free list
        free_blocks.remove(&block_offset);

        let ptr = self.base_address + aligned_offset;

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.allocation_count += 1;
            stats.allocated_bytes += alloc_size;
            stats.free_bytes -= alloc_size;
            stats.free_block_count = free_blocks.len();
            stats.peak_allocated_bytes = stats.peak_allocated_bytes.max(stats.allocated_bytes);

            if stats.allocation_count > 0 {
                stats.avg_allocation_size = stats.allocated_bytes / stats.allocation_count;
            }
        }

        Ok(MemoryAllocation {
            block_id,
            ptr,
            size: alloc_size,
            device_id: self.device_id,
        })
    }

    /// Free an allocation
    pub fn free(&self, allocation: &MemoryAllocation) -> Result<()> {
        let mut blocks = self.blocks.write();
        let mut free_blocks = self.free_blocks.write();

        let block = blocks
            .get_mut(&allocation.block_id)
            .ok_or_else(|| DistributedError::runtime("Block not found"))?;

        if block.state != BlockState::Allocated {
            return Err(DistributedError::runtime("Block is not allocated"));
        }

        let block_size = block.size;
        let block_offset = block.offset;

        // Try to cache the block
        if self.config.cache_freed_blocks {
            let mut cache = self.size_cache.lock();
            if cache.put(block_size, allocation.block_id) {
                block.state = BlockState::Cached;
                block.allocated_at = None;
                block.tag = None;

                // Update stats
                let mut stats = self.stats.write();
                stats.deallocation_count += 1;
                stats.allocated_bytes -= block_size;
                stats.free_bytes += block_size;

                return Ok(());
            }
        }

        // Mark as free
        block.state = BlockState::Free;
        block.allocated_at = None;
        block.tag = None;

        // Add to free list
        free_blocks.insert(block_offset, allocation.block_id);

        // Update stats
        {
            let mut stats = self.stats.write();
            stats.deallocation_count += 1;
            stats.allocated_bytes -= block_size;
            stats.free_bytes += block_size;
            stats.free_block_count = free_blocks.len();
        }

        // Try to coalesce with neighbors
        drop(blocks);
        drop(free_blocks);
        self.coalesce_neighbors(allocation.block_id)?;

        // Check if defrag is needed
        if self.config.auto_defrag {
            let stats = self.stats.read();
            if stats.fragmentation_ratio > self.config.defrag_threshold {
                drop(stats);
                self.defragment()?;
            }
        }

        Ok(())
    }

    /// Coalesce free blocks with neighbors
    fn coalesce_neighbors(&self, block_id: BlockId) -> Result<()> {
        let mut blocks = self.blocks.write();
        let mut free_blocks = self.free_blocks.write();

        let block = match blocks.get(&block_id) {
            Some(b) if b.state == BlockState::Free => b.clone(),
            _ => return Ok(()),
        };

        // Find and merge with next block
        let next_offset = block.offset + block.size;
        if let Some(&next_id) = free_blocks.get(&next_offset) {
            if let Some(next_block) = blocks.get(&next_id) {
                if next_block.state == BlockState::Free {
                    let merged_size = block.size + next_block.size;

                    // Remove next block
                    free_blocks.remove(&next_offset);
                    blocks.remove(&next_id);

                    // Update current block
                    if let Some(current) = blocks.get_mut(&block_id) {
                        current.size = merged_size;
                    }
                }
            }
        }

        // Update fragmentation
        self.update_fragmentation(&blocks, &free_blocks);

        Ok(())
    }

    /// Update fragmentation ratio
    fn update_fragmentation(
        &self,
        blocks: &HashMap<BlockId, MemoryBlock>,
        free_blocks: &BTreeMap<usize, BlockId>,
    ) {
        if free_blocks.is_empty() {
            let mut stats = self.stats.write();
            stats.fragmentation_ratio = 0.0;
            return;
        }

        // Calculate fragmentation as ratio of free blocks to total free space
        let mut total_free = 0;
        let mut largest_free = 0;

        for &block_id in free_blocks.values() {
            if let Some(block) = blocks.get(&block_id) {
                total_free += block.size;
                largest_free = largest_free.max(block.size);
            }
        }

        let mut stats = self.stats.write();
        if total_free > 0 {
            // Fragmentation = 1 - (largest_free / total_free)
            stats.fragmentation_ratio = 1.0 - (largest_free as f32 / total_free as f32);
        } else {
            stats.fragmentation_ratio = 0.0;
        }
    }

    /// Defragment the pool
    pub fn defragment(&self) -> Result<()> {
        // In a real implementation, this would:
        // 1. Wait for all allocations to complete
        // 2. Copy data to compact free space
        // 3. Update block offsets
        // For simulation, we just clear the cache and update stats

        let mut cache = self.size_cache.lock();
        cache.clear();

        let mut stats = self.stats.write();
        stats.defrag_count += 1;

        Ok(())
    }

    /// Align size to minimum block size
    fn align_size(&self, size: usize) -> usize {
        let min = self.config.min_block_size;
        ((size + min - 1) / min) * min
    }

    /// Get pool statistics
    pub fn stats(&self) -> MemoryPoolStats {
        self.stats.read().clone()
    }

    /// Get device ID
    pub fn device_id(&self) -> i32 {
        self.device_id
    }

    /// Get current pool size
    pub fn current_size(&self) -> usize {
        self.current_size.load(Ordering::SeqCst)
    }

    /// Clear the pool
    pub fn clear(&self) -> Result<()> {
        let mut blocks = self.blocks.write();
        let mut free_blocks = self.free_blocks.write();
        let mut cache = self.size_cache.lock();

        blocks.clear();
        free_blocks.clear();
        cache.clear();

        // Reinitialize with single free block
        let size = self.current_size.load(Ordering::SeqCst);
        let block = MemoryBlock::new(0, 0, size, self.device_id);
        blocks.insert(0, block);
        free_blocks.insert(0, 0);

        // Reset stats
        let mut stats = self.stats.write();
        *stats = MemoryPoolStats {
            total_size_bytes: size,
            free_bytes: size,
            free_block_count: 1,
            ..Default::default()
        };

        Ok(())
    }
}

// =============================================================================
// Memory Allocation Handle
// =============================================================================

/// Handle to an allocated memory region
#[derive(Debug, Clone)]
pub struct MemoryAllocation {
    /// Block ID in the pool
    pub block_id: BlockId,
    /// Pointer to allocated memory
    pub ptr: usize,
    /// Size in bytes
    pub size: usize,
    /// Device ID
    pub device_id: i32,
}

impl MemoryAllocation {
    /// Get the memory address
    pub fn address(&self) -> usize {
        self.ptr
    }

    /// Get size in bytes
    pub fn size_bytes(&self) -> usize {
        self.size
    }
}

// =============================================================================
// Thread-Safe Wrapper
// =============================================================================

/// Thread-safe shared memory pool
pub type SharedMemoryPool = Arc<MemoryPool>;

/// Create a shared memory pool
pub fn shared_memory_pool(config: MemoryPoolConfig, device_id: i32) -> Result<SharedMemoryPool> {
    Ok(Arc::new(MemoryPool::new(config, device_id)?))
}

// =============================================================================
// Multi-Device Pool Manager
// =============================================================================

/// Manages memory pools across multiple devices
pub struct MultiDevicePoolManager {
    /// Pools per device
    pools: HashMap<i32, SharedMemoryPool>,
    /// Configuration
    config: MemoryPoolConfig,
}

impl MultiDevicePoolManager {
    /// Create a new multi-device pool manager
    pub fn new(config: MemoryPoolConfig) -> Self {
        Self {
            pools: HashMap::new(),
            config,
        }
    }

    /// Get or create pool for a device
    pub fn get_pool(&mut self, device_id: i32) -> Result<SharedMemoryPool> {
        if let Some(pool) = self.pools.get(&device_id) {
            return Ok(pool.clone());
        }

        let pool = shared_memory_pool(self.config.clone(), device_id)?;
        self.pools.insert(device_id, pool.clone());
        Ok(pool)
    }

    /// Get statistics for all devices
    pub fn all_stats(&self) -> HashMap<i32, MemoryPoolStats> {
        self.pools
            .iter()
            .map(|(&id, pool)| (id, pool.stats()))
            .collect()
    }

    /// Clear all pools
    pub fn clear_all(&self) -> Result<()> {
        for pool in self.pools.values() {
            pool.clear()?;
        }
        Ok(())
    }
}

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

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

    #[test]
    fn test_memory_pool_config_default() {
        let config = MemoryPoolConfig::default();
        assert_eq!(config.initial_size_bytes, 256 * 1024 * 1024);
        assert_eq!(config.strategy, AllocationStrategy::BestFit);
        assert_eq!(config.alignment, 256);
    }

    #[test]
    fn test_memory_block_creation() {
        let block = MemoryBlock::new(1, 1024, 4096, 0);
        assert_eq!(block.id, 1);
        assert_eq!(block.offset, 1024);
        assert_eq!(block.size, 4096);
        assert_eq!(block.state, BlockState::Free);
    }

    #[test]
    fn test_block_can_satisfy() {
        let block = MemoryBlock::new(1, 0, 4096, 0);
        assert!(block.can_satisfy(1024, 256));
        assert!(block.can_satisfy(4096, 256));
        assert!(!block.can_satisfy(8192, 256));
    }

    #[test]
    fn test_memory_pool_creation() {
        let config = MemoryPoolConfig::default();
        let pool = MemoryPool::new(config, 0).unwrap();

        assert_eq!(pool.device_id(), 0);
        assert_eq!(pool.current_size(), 256 * 1024 * 1024);
    }

    #[test]
    fn test_simple_allocation() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024, // 1MB
            ..Default::default()
        };
        let pool = MemoryPool::new(config, 0).unwrap();

        let alloc = pool.allocate(4096).unwrap();
        assert!(alloc.ptr > 0);
        assert!(alloc.size >= 4096);

        let stats = pool.stats();
        assert_eq!(stats.allocation_count, 1);
        assert!(stats.allocated_bytes >= 4096);
    }

    #[test]
    fn test_allocation_and_free() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024,
            ..Default::default()
        };
        let pool = MemoryPool::new(config, 0).unwrap();

        let alloc = pool.allocate(4096).unwrap();
        let stats_before = pool.stats();

        pool.free(&alloc).unwrap();
        let stats_after = pool.stats();

        assert_eq!(stats_after.deallocation_count, 1);
        assert!(stats_after.allocated_bytes < stats_before.allocated_bytes);
    }

    #[test]
    fn test_multiple_allocations() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024,
            ..Default::default()
        };
        let pool = MemoryPool::new(config, 0).unwrap();

        let alloc1 = pool.allocate(1024).unwrap();
        let alloc2 = pool.allocate(2048).unwrap();
        let alloc3 = pool.allocate(4096).unwrap();

        assert_ne!(alloc1.ptr, alloc2.ptr);
        assert_ne!(alloc2.ptr, alloc3.ptr);

        let stats = pool.stats();
        assert_eq!(stats.allocation_count, 3);
    }

    #[test]
    fn test_size_class_cache() {
        let mut cache = SizeClassCache::new(10, 512);

        cache.put(1024, 1);
        cache.put(1024, 2);
        cache.put(2048, 3);

        assert_eq!(cache.get(1024), Some(1));
        assert_eq!(cache.get(1024), Some(2));
        assert_eq!(cache.get(1024), None);
        assert_eq!(cache.get(2048), Some(3));
    }

    #[test]
    fn test_cache_hit() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024,
            cache_freed_blocks: true,
            ..Default::default()
        };
        let pool = MemoryPool::new(config, 0).unwrap();

        // Allocate and free
        let alloc1 = pool.allocate(4096).unwrap();
        pool.free(&alloc1).unwrap();

        // Allocate same size - should hit cache
        let _alloc2 = pool.allocate(4096).unwrap();

        let stats = pool.stats();
        assert!(stats.cache_hits >= 1);
    }

    #[test]
    fn test_allocation_with_tag() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024,
            ..Default::default()
        };
        let pool = MemoryPool::new(config, 0).unwrap();

        let alloc = pool
            .allocate_with_tag(4096, Some("gradient".to_string()))
            .unwrap();
        assert!(alloc.size >= 4096);
    }

    #[test]
    fn test_pool_clear() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024,
            ..Default::default()
        };
        let pool = MemoryPool::new(config, 0).unwrap();

        let _ = pool.allocate(4096).unwrap();
        let _ = pool.allocate(4096).unwrap();

        pool.clear().unwrap();

        let stats = pool.stats();
        assert_eq!(stats.allocation_count, 0);
        assert_eq!(stats.free_block_count, 1);
    }

    #[test]
    fn test_multi_device_manager() {
        let config = MemoryPoolConfig {
            initial_size_bytes: 1024 * 1024,
            ..Default::default()
        };
        let mut manager = MultiDevicePoolManager::new(config);

        let pool0 = manager.get_pool(0).unwrap();
        let pool1 = manager.get_pool(1).unwrap();

        assert_eq!(pool0.device_id(), 0);
        assert_eq!(pool1.device_id(), 1);

        let all_stats = manager.all_stats();
        assert_eq!(all_stats.len(), 2);
    }

    #[test]
    fn test_memory_pool_stats() {
        let stats = MemoryPoolStats {
            total_size_bytes: 1000,
            allocated_bytes: 500,
            cache_hits: 10,
            cache_misses: 5,
            ..Default::default()
        };

        assert!((stats.utilization() - 0.5).abs() < 0.01);
        assert!((stats.cache_hit_ratio() - 0.666).abs() < 0.01);
    }

    #[test]
    fn test_allocation_strategy() {
        assert_eq!(AllocationStrategy::BestFit, AllocationStrategy::BestFit);
        assert_ne!(AllocationStrategy::BestFit, AllocationStrategy::FirstFit);
    }
}