ruviz 0.6.0

use crate::core::types::Point2f;
use std::alloc::{Layout, alloc, dealloc};
use std::sync::{Arc, Mutex, OnceLock};

/// High-performance memory management system for plotting operations
///
/// Provides object pooling, pre-allocation strategies, and memory-efficient
/// data structures optimized for repeated plotting operations.
///
/// # Locking policy
///
/// `buffer_pools`, `stats`, and the nested `BlockPool` mutex are never held
/// simultaneously. Operations spanning them use separate critical sections.
#[derive(Debug)]
pub struct MemoryManager {
    /// Buffer pools for different data types
    buffer_pools: Arc<Mutex<BufferPools>>,
    /// Memory usage statistics
    stats: Arc<Mutex<MemoryStats>>,
    /// Configuration parameters
    config: MemoryConfig,
    #[cfg(test)]
    lock_test_hooks: Arc<Mutex<LockTestHooks>>,
}

#[cfg(test)]
#[derive(Debug, Default)]
struct LockTestHooks {
    after_buffer_pool_access: Option<(std::sync::mpsc::SyncSender<()>, Arc<std::sync::Barrier>)>,
    after_stats_snapshot: Option<(std::sync::mpsc::SyncSender<()>, Arc<std::sync::Barrier>)>,
}

/// Buffer pools for different primitive types
#[derive(Debug)]
struct BufferPools {
    /// Pool for f32 vectors (coordinates, colors, etc.)
    f32_buffers: BufferPool<f32>,
    /// Pool for f64 vectors (input data)
    f64_buffers: BufferPool<f64>,
    /// Pool for u8 vectors (pixel data, colors)
    u8_buffers: BufferPool<u8>,
    /// Pool for u32 vectors (indices, counts)
    u32_buffers: BufferPool<u32>,
    /// Pool for Point2f vectors
    point_buffers: BufferPool<Point2f>,
    /// Pool for large allocation blocks
    block_pool: Arc<Mutex<BlockPool>>,
}

/// Generic buffer pool for reusable vectors
#[derive(Debug)]
struct BufferPool<T> {
    /// Available buffers sorted by capacity
    available: Vec<Vec<T>>,
    /// Currently allocated buffer count
    allocated_count: usize,
    /// Maximum pool size to prevent unbounded growth
    max_pool_size: usize,
    /// Minimum buffer capacity to pool
    min_capacity: usize,
}

/// Pool for large memory blocks
#[derive(Debug)]
struct BlockPool {
    /// Available memory blocks
    blocks: Vec<MemoryBlock>,
    /// Block allocation statistics
    stats: BlockStats,
}

/// Memory block for large allocations
#[derive(Debug, Clone)]
struct MemoryBlock {
    /// Pointer to allocated memory
    ptr: *mut u8,
    /// Size of the allocation
    size: usize,
    /// Layout used for allocation
    layout: Layout,
}

// SAFETY: `ptr` is treated as an owning allocation handle and all access/deallocation
// is synchronized through `Mutex<BlockPool>` or RAII `ManagedBlock` drop paths.
unsafe impl Send for MemoryBlock {}

/// Memory usage statistics
#[derive(Debug, Clone)]
pub struct MemoryStats {
    /// Total bytes allocated
    pub total_allocated: usize,
    /// Total bytes deallocated
    pub total_deallocated: usize,
    /// Current memory usage
    pub current_usage: usize,
    /// Peak memory usage
    pub peak_usage: usize,
    /// Number of active allocations
    pub active_allocations: usize,
    /// Pool hit rate (reused buffers / total requests)
    pub pool_hit_rate: f32,
    /// Buffer pool statistics
    pub pool_stats: PoolStats,
}

/// Buffer pool statistics
#[derive(Debug, Clone)]
pub struct PoolStats {
    pub f32_pool_size: usize,
    pub f64_pool_size: usize,
    pub u8_pool_size: usize,
    pub u32_pool_size: usize,
    pub point_pool_size: usize,
    pub block_pool_size: usize,
    pub total_pool_memory: usize,
}

/// Block allocation statistics
#[derive(Debug, Clone)]
struct BlockStats {
    total_blocks_allocated: usize,
    total_blocks_reused: usize,
    peak_block_count: usize,
}

/// Memory configuration
#[derive(Debug, Clone)]
pub struct MemoryConfig {
    /// Enable buffer pooling
    pub enable_pooling: bool,
    /// Maximum pool size for each buffer type
    pub max_pool_size: usize,
    /// Minimum capacity to consider for pooling
    pub min_pool_capacity: usize,
    /// Pre-allocate buffers for common sizes
    pub pre_allocate_common_sizes: bool,
    /// Enable memory usage tracking
    pub track_usage: bool,
    /// Large allocation threshold (bytes)
    pub large_alloc_threshold: usize,
    /// Maximum total pool memory (bytes)
    pub max_pool_memory: usize,
}

/// Managed buffer that returns to pool when dropped
pub struct ManagedBuffer<T> {
    buffer: Option<Vec<T>>,
    recycler: Option<Arc<dyn Fn(Vec<T>) + Send + Sync>>,
    stats: Arc<Mutex<MemoryStats>>,
}

impl<T: std::fmt::Debug> std::fmt::Debug for ManagedBuffer<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ManagedBuffer")
            .field("buffer", &self.buffer)
            .finish_non_exhaustive()
    }
}

impl Default for MemoryManager {
    fn default() -> Self {
        Self::new()
    }
}

impl MemoryManager {
    /// Create new memory manager with default configuration
    pub fn new() -> Self {
        Self::with_config(MemoryConfig::default())
    }

    /// Create memory manager with custom configuration
    pub fn with_config(config: MemoryConfig) -> Self {
        let buffer_pools = BufferPools::new(&config);

        Self {
            buffer_pools: Arc::new(Mutex::new(buffer_pools)),
            stats: Arc::new(Mutex::new(MemoryStats::new())),
            config,
            #[cfg(test)]
            lock_test_hooks: Arc::new(Mutex::new(LockTestHooks::default())),
        }
    }

    /// Get or create a managed f32 buffer
    pub fn get_f32_buffer(&self, min_capacity: usize) -> ManagedBuffer<f32> {
        if !self.config.enable_pooling {
            return ManagedBuffer::new_unmanaged(Vec::with_capacity(min_capacity));
        }

        let (buffer, reused) = {
            let mut pools = self.buffer_pools.lock().unwrap();
            pools.f32_buffers.get_buffer(min_capacity)
        };
        #[cfg(test)]
        self.wait_after_buffer_pool_access_for_test();

        {
            let mut stats = self.stats.lock().unwrap();
            stats.active_allocations += 1;
            if reused {
                stats.update_pool_hit();
            }
        }

        let pool_arc = self.buffer_pools.clone();
        ManagedBuffer {
            buffer: Some(buffer),
            recycler: Some(Arc::new(move |buffer: Vec<f32>| {
                if let Ok(mut pools) = pool_arc.lock() {
                    pools.f32_buffers.return_buffer(buffer);
                }
            })),
            stats: self.stats.clone(),
        }
    }

    /// Get or create a managed f64 buffer
    pub fn get_f64_buffer(&self, min_capacity: usize) -> ManagedBuffer<f64> {
        if !self.config.enable_pooling {
            return ManagedBuffer::new_unmanaged(Vec::with_capacity(min_capacity));
        }

        let (buffer, reused) = {
            let mut pools = self.buffer_pools.lock().unwrap();
            pools.f64_buffers.get_buffer(min_capacity)
        };
        #[cfg(test)]
        self.wait_after_buffer_pool_access_for_test();

        {
            let mut stats = self.stats.lock().unwrap();
            stats.active_allocations += 1;
            if reused {
                stats.update_pool_hit();
            }
        }

        let pool_arc = self.buffer_pools.clone();
        ManagedBuffer {
            buffer: Some(buffer),
            recycler: Some(Arc::new(move |buffer: Vec<f64>| {
                if let Ok(mut pools) = pool_arc.lock() {
                    pools.f64_buffers.return_buffer(buffer);
                }
            })),
            stats: self.stats.clone(),
        }
    }

    /// Get or create a managed Point2f buffer
    pub fn get_point_buffer(&self, min_capacity: usize) -> ManagedBuffer<Point2f> {
        if !self.config.enable_pooling {
            return ManagedBuffer::new_unmanaged(Vec::with_capacity(min_capacity));
        }

        let (buffer, reused) = {
            let mut pools = self.buffer_pools.lock().unwrap();
            pools.point_buffers.get_buffer(min_capacity)
        };
        #[cfg(test)]
        self.wait_after_buffer_pool_access_for_test();

        {
            let mut stats = self.stats.lock().unwrap();
            stats.active_allocations += 1;
            if reused {
                stats.update_pool_hit();
            }
        }

        let pool_arc = self.buffer_pools.clone();
        ManagedBuffer {
            buffer: Some(buffer),
            recycler: Some(Arc::new(move |buffer: Vec<Point2f>| {
                if let Ok(mut pools) = pool_arc.lock() {
                    pools.point_buffers.return_buffer(buffer);
                }
            })),
            stats: self.stats.clone(),
        }
    }

    /// Get or create a managed u8 buffer for pixel data
    pub fn get_u8_buffer(&self, min_capacity: usize) -> ManagedBuffer<u8> {
        if !self.config.enable_pooling {
            return ManagedBuffer::new_unmanaged(Vec::with_capacity(min_capacity));
        }

        let (buffer, reused) = {
            let mut pools = self.buffer_pools.lock().unwrap();
            pools.u8_buffers.get_buffer(min_capacity)
        };
        #[cfg(test)]
        self.wait_after_buffer_pool_access_for_test();

        {
            let mut stats = self.stats.lock().unwrap();
            stats.active_allocations += 1;
            if reused {
                stats.update_pool_hit();
            }
        }

        let pool_arc = self.buffer_pools.clone();
        ManagedBuffer {
            buffer: Some(buffer),
            recycler: Some(Arc::new(move |buffer: Vec<u8>| {
                if let Ok(mut pools) = pool_arc.lock() {
                    pools.u8_buffers.return_buffer(buffer);
                }
            })),
            stats: self.stats.clone(),
        }
    }

    /// Allocate a large memory block
    pub fn allocate_block(
        &self,
        size: usize,
        alignment: usize,
    ) -> Result<ManagedBlock, MemoryError> {
        if size >= self.config.large_alloc_threshold {
            let pool_arc = {
                let pools = self.buffer_pools.lock().unwrap();
                pools.block_pool.clone()
            };
            let mut pool = pool_arc.lock().unwrap();
            return pool.allocate_block(size, alignment, pool_arc.clone());
        }

        // For smaller allocations, use regular allocation
        let layout =
            Layout::from_size_align(size, alignment).map_err(|_| MemoryError::InvalidLayout)?;

        // SAFETY: `layout` is validated above, and we check null before storing `ptr`.
        unsafe {
            let ptr = alloc(layout);
            if ptr.is_null() {
                return Err(MemoryError::AllocationFailed);
            }

            Ok(ManagedBlock {
                ptr,
                size,
                layout,
                pool: None,
            })
        }
    }

    /// Pre-allocate buffers for common sizes
    pub fn pre_allocate_common_sizes(&self) {
        if !self.config.pre_allocate_common_sizes {
            return;
        }

        let common_sizes = [
            100,     // Small datasets
            1_000,   // Medium datasets
            10_000,  // Large datasets
            100_000, // Very large datasets
        ];

        let mut pools = self.buffer_pools.lock().unwrap();

        for &size in &common_sizes {
            // Pre-allocate multiple buffers for each common size
            for _ in 0..3 {
                pools.f32_buffers.pre_allocate(size);
                pools.f64_buffers.pre_allocate(size);
                pools.point_buffers.pre_allocate(size);
            }

            // Pre-allocate pixel buffers (typically larger)
            let pixel_size = size * 4; // RGBA
            pools.u8_buffers.pre_allocate(pixel_size);
        }
    }

    /// Get current memory statistics.
    ///
    /// This is a best-effort telemetry snapshot, not an atomic transaction across
    /// allocation counters and pools. Each component is internally consistent,
    /// but concurrent allocation or recycling may occur between the snapshots.
    pub fn get_stats(&self) -> MemoryStats {
        let mut stats_copy = {
            let stats = self.stats.lock().unwrap();
            stats.clone()
        };
        #[cfg(test)]
        self.wait_after_stats_snapshot_for_test();

        let (mut pool_stats, block_pool) = {
            let pools = self.buffer_pools.lock().unwrap();
            (pools.get_buffer_pool_stats(), pools.block_pool.clone())
        };
        let block_pool_stats = block_pool
            .lock()
            .map(|pool| (pool.blocks.len(), pool.memory_usage()))
            .unwrap_or((0, 0));
        pool_stats.block_pool_size = block_pool_stats.0;
        pool_stats.total_pool_memory += block_pool_stats.1;
        stats_copy.pool_stats = pool_stats;

        stats_copy
    }

    /// Clear all pools and reset statistics.
    ///
    /// Clearing is intentionally split into independent critical sections, so
    /// concurrent snapshots may briefly observe either side of the reset.
    pub fn clear(&self) {
        let block_pool = {
            let mut pools = self.buffer_pools.lock().unwrap();
            pools.clear_buffers();
            pools.block_pool.clone()
        };
        if let Ok(mut block_pool) = block_pool.lock() {
            block_pool.clear();
        }
        self.stats.lock().unwrap().reset();
    }

    #[cfg(test)]
    fn set_after_buffer_pool_access_hook(
        &self,
        sender: std::sync::mpsc::SyncSender<()>,
        barrier: Arc<std::sync::Barrier>,
    ) {
        self.lock_test_hooks
            .lock()
            .unwrap()
            .after_buffer_pool_access = Some((sender, barrier));
    }

    #[cfg(test)]
    fn wait_after_buffer_pool_access_for_test(&self) {
        let hook = self
            .lock_test_hooks
            .lock()
            .unwrap()
            .after_buffer_pool_access
            .take();
        if let Some((sender, barrier)) = hook {
            sender.send(()).unwrap();
            barrier.wait();
        }
    }

    #[cfg(test)]
    fn set_after_stats_snapshot_hook(
        &self,
        sender: std::sync::mpsc::SyncSender<()>,
        barrier: Arc<std::sync::Barrier>,
    ) {
        self.lock_test_hooks.lock().unwrap().after_stats_snapshot = Some((sender, barrier));
    }

    #[cfg(test)]
    fn wait_after_stats_snapshot_for_test(&self) {
        let hook = self
            .lock_test_hooks
            .lock()
            .unwrap()
            .after_stats_snapshot
            .take();
        if let Some((sender, barrier)) = hook {
            sender.send(()).unwrap();
            barrier.wait();
        }
    }

    /// Get memory configuration
    pub fn config(&self) -> &MemoryConfig {
        &self.config
    }
}

impl<T> BufferPool<T> {
    fn new(max_pool_size: usize, min_capacity: usize) -> Self {
        Self {
            available: Vec::new(),
            allocated_count: 0,
            max_pool_size,
            min_capacity,
        }
    }

    fn get_buffer(&mut self, min_capacity: usize) -> (Vec<T>, bool) {
        // Try to find a suitable buffer in the pool
        if min_capacity >= self.min_capacity
            && let Some(pos) = self
                .available
                .iter()
                .position(|buf| buf.capacity() >= min_capacity)
        {
            let mut buffer = self.available.swap_remove(pos);
            buffer.clear();
            return (buffer, true);
        }

        // No suitable buffer found, allocate new one
        self.allocated_count += 1;
        (Vec::with_capacity(min_capacity), false)
    }

    fn return_buffer(&mut self, mut buffer: Vec<T>) {
        buffer.clear();

        // Only keep buffer if pool isn't full and capacity is worth pooling
        if self.available.len() < self.max_pool_size && buffer.capacity() >= self.min_capacity {
            // Insert in sorted order by capacity for better reuse
            let insert_pos = self
                .available
                .binary_search_by_key(&buffer.capacity(), |buf| buf.capacity())
                .unwrap_or_else(|pos| pos);

            self.available.insert(insert_pos, buffer);
        }

        self.allocated_count = self.allocated_count.saturating_sub(1);
    }

    fn pre_allocate(&mut self, capacity: usize) {
        if self.available.len() < self.max_pool_size {
            let buffer = Vec::with_capacity(capacity);
            self.available.push(buffer);
        }
    }

    fn clear(&mut self) {
        self.available.clear();
        self.allocated_count = 0;
    }

    fn memory_usage(&self) -> usize {
        self.available
            .iter()
            .map(|buf| buf.capacity() * std::mem::size_of::<T>())
            .sum()
    }
}

impl BufferPools {
    fn new(config: &MemoryConfig) -> Self {
        Self {
            f32_buffers: BufferPool::new(config.max_pool_size, config.min_pool_capacity),
            f64_buffers: BufferPool::new(config.max_pool_size, config.min_pool_capacity),
            u8_buffers: BufferPool::new(config.max_pool_size, config.min_pool_capacity),
            u32_buffers: BufferPool::new(config.max_pool_size, config.min_pool_capacity),
            point_buffers: BufferPool::new(config.max_pool_size, config.min_pool_capacity),
            block_pool: Arc::new(Mutex::new(BlockPool::new())),
        }
    }

    fn clear_buffers(&mut self) {
        self.f32_buffers.clear();
        self.f64_buffers.clear();
        self.u8_buffers.clear();
        self.u32_buffers.clear();
        self.point_buffers.clear();
    }

    fn get_buffer_pool_stats(&self) -> PoolStats {
        PoolStats {
            f32_pool_size: self.f32_buffers.available.len(),
            f64_pool_size: self.f64_buffers.available.len(),
            u8_pool_size: self.u8_buffers.available.len(),
            u32_pool_size: self.u32_buffers.available.len(),
            point_pool_size: self.point_buffers.available.len(),
            block_pool_size: 0,
            total_pool_memory: self.f32_buffers.memory_usage()
                + self.f64_buffers.memory_usage()
                + self.u8_buffers.memory_usage()
                + self.u32_buffers.memory_usage()
                + self.point_buffers.memory_usage(),
        }
    }
}

impl BlockPool {
    fn new() -> Self {
        Self {
            blocks: Vec::new(),
            stats: BlockStats {
                total_blocks_allocated: 0,
                total_blocks_reused: 0,
                peak_block_count: 0,
            },
        }
    }

    fn allocate_block(
        &mut self,
        size: usize,
        alignment: usize,
        pool: Arc<Mutex<BlockPool>>,
    ) -> Result<ManagedBlock, MemoryError> {
        // Try to find a suitable existing block
        if let Some(pos) = self.blocks.iter().position(|block| block.size >= size) {
            let block = self.blocks.swap_remove(pos);
            self.stats.total_blocks_reused += 1;

            return Ok(ManagedBlock {
                ptr: block.ptr,
                size: block.size,
                layout: block.layout,
                pool: Some(pool),
            });
        }

        // Allocate new block
        let layout =
            Layout::from_size_align(size, alignment).map_err(|_| MemoryError::InvalidLayout)?;

        // SAFETY: `layout` is validated above, and `ptr` is checked for null before use.
        unsafe {
            let ptr = alloc(layout);
            if ptr.is_null() {
                return Err(MemoryError::AllocationFailed);
            }

            self.stats.total_blocks_allocated += 1;
            self.stats.peak_block_count = self.stats.peak_block_count.max(self.blocks.len() + 1);

            Ok(ManagedBlock {
                ptr,
                size,
                layout,
                pool: Some(pool),
            })
        }
    }

    fn return_block(&mut self, block: MemoryBlock) {
        self.blocks.push(block);
    }

    fn clear(&mut self) {
        for block in self.blocks.drain(..) {
            // SAFETY: each block was allocated with this exact layout and is deallocated once.
            unsafe {
                dealloc(block.ptr, block.layout);
            }
        }
        self.stats = BlockStats {
            total_blocks_allocated: 0,
            total_blocks_reused: 0,
            peak_block_count: 0,
        };
    }

    fn memory_usage(&self) -> usize {
        self.blocks.iter().map(|block| block.size).sum()
    }
}

impl MemoryStats {
    fn new() -> Self {
        Self {
            total_allocated: 0,
            total_deallocated: 0,
            current_usage: 0,
            peak_usage: 0,
            active_allocations: 0,
            pool_hit_rate: 0.0,
            pool_stats: PoolStats {
                f32_pool_size: 0,
                f64_pool_size: 0,
                u8_pool_size: 0,
                u32_pool_size: 0,
                point_pool_size: 0,
                block_pool_size: 0,
                total_pool_memory: 0,
            },
        }
    }

    fn update_pool_hit(&mut self) {
        // Simple moving average for pool hit rate
        self.pool_hit_rate = self.pool_hit_rate * 0.9 + 0.1;
    }

    fn reset(&mut self) {
        *self = Self::new();
    }
}

impl Default for MemoryConfig {
    fn default() -> Self {
        Self {
            enable_pooling: true,
            max_pool_size: 10,
            min_pool_capacity: 100,
            pre_allocate_common_sizes: true,
            track_usage: true,
            large_alloc_threshold: 1024 * 1024, // 1MB
            max_pool_memory: 100 * 1024 * 1024, // 100MB
        }
    }
}

impl<T> ManagedBuffer<T> {
    fn new_unmanaged(buffer: Vec<T>) -> Self {
        Self {
            buffer: Some(buffer),
            recycler: None,
            stats: Arc::new(Mutex::new(MemoryStats::new())),
        }
    }

    /// Get mutable access to the underlying buffer
    pub fn get_mut(&mut self) -> &mut Vec<T> {
        self.buffer.as_mut().unwrap()
    }

    /// Get read access to the underlying buffer
    pub fn get(&self) -> &Vec<T> {
        self.buffer.as_ref().unwrap()
    }

    /// Take ownership of the buffer (consumes the managed buffer)
    ///
    /// The storage escapes the pool — the caller owns it now, so it cannot be
    /// recycled — but the allocation still stops being *managed*, so it is
    /// accounted for exactly as a drop is.
    pub fn into_inner(mut self) -> Vec<T> {
        self.release(Recycle::No).unwrap_or_default()
    }

    /// Release the managed buffer, optionally handing the storage back to the
    /// pool, and balance `active_allocations` once.
    ///
    /// Both exits from a `ManagedBuffer` — `into_inner` and `Drop` — funnel
    /// through here so the accounting cannot diverge between them. It used to:
    /// `into_inner` took the `Option` itself, which left `Drop` with nothing to
    /// do, so every buffer handed out through `into_inner` was counted as live
    /// forever and `MemoryStats::active_allocations` grew without bound.
    ///
    /// Returns the buffer when it was not recycled.
    fn release(&mut self, recycle: Recycle) -> Option<Vec<T>> {
        let buffer = self.buffer.take()?;

        // Take the pool lock (inside `recycler`) before the stats lock, matching
        // the order used when the buffer was handed out.
        let escaped = match (recycle, &self.recycler) {
            (Recycle::Yes, Some(recycler)) => {
                recycler(buffer);
                None
            }
            _ => Some(buffer),
        };

        if let Ok(mut stats) = self.stats.lock() {
            stats.active_allocations = stats.active_allocations.saturating_sub(1);
        }

        escaped
    }
}

/// Whether a released [`ManagedBuffer`] hands its storage back to the pool.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Recycle {
    Yes,
    No,
}

impl<T> Drop for ManagedBuffer<T> {
    fn drop(&mut self) {
        let _ = self.release(Recycle::Yes);
    }
}

/// Managed memory block
pub struct ManagedBlock {
    ptr: *mut u8,
    size: usize,
    layout: Layout,
    pool: Option<Arc<Mutex<BlockPool>>>,
}

// SAFETY: `ManagedBlock` only carries ownership metadata for an allocation and does not
// provide unsynchronized shared mutation of pointee memory.
unsafe impl Send for ManagedBlock {}

impl Drop for ManagedBlock {
    fn drop(&mut self) {
        if let Some(pool_arc) = &self.pool {
            let mut pool = pool_arc.lock().unwrap();
            pool.return_block(MemoryBlock {
                ptr: self.ptr,
                size: self.size,
                layout: self.layout,
            });
        } else {
            // SAFETY: this block was allocated by `alloc(layout)` and not returned to pool.
            unsafe {
                dealloc(self.ptr, self.layout);
            }
        }
    }
}

/// Memory allocation errors
#[derive(Debug, Clone)]
pub enum MemoryError {
    AllocationFailed,
    InvalidLayout,
    PoolExhausted,
}

impl std::fmt::Display for MemoryError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MemoryError::AllocationFailed => write!(f, "Memory allocation failed"),
            MemoryError::InvalidLayout => write!(f, "Invalid memory layout"),
            MemoryError::PoolExhausted => write!(f, "Memory pool exhausted"),
        }
    }
}

impl std::error::Error for MemoryError {}

/// Global memory manager instance
static GLOBAL_MEMORY_MANAGER: OnceLock<MemoryManager> = OnceLock::new();

/// Get the global memory manager instance
pub fn get_memory_manager() -> &'static MemoryManager {
    GLOBAL_MEMORY_MANAGER.get_or_init(|| {
        let config = MemoryConfig::default();
        let manager = MemoryManager::with_config(config);
        manager.pre_allocate_common_sizes();
        manager
    })
}

/// Initialize global memory manager with custom configuration
/// Note: This only works if called before first access to get_memory_manager
pub fn initialize_memory_manager(config: MemoryConfig) -> Result<(), MemoryError> {
    let manager = MemoryManager::with_config(config);
    manager.pre_allocate_common_sizes();

    GLOBAL_MEMORY_MANAGER
        .set(manager)
        .map_err(|_| MemoryError::InvalidLayout) // Already initialized
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Barrier, mpsc};
    use std::thread;
    use std::time::Duration;

    const LOCK_TEST_TIMEOUT: Duration = Duration::from_secs(5);

    #[test]
    fn test_memory_manager_creation() {
        let manager = MemoryManager::new();
        let stats = manager.get_stats();

        assert_eq!(stats.active_allocations, 0);
        assert_eq!(stats.current_usage, 0);
    }

    #[test]
    fn test_buffer_pooling() {
        let manager = MemoryManager::new();

        // Get a buffer
        {
            let _buffer = manager.get_f32_buffer(1000);
            let stats = manager.get_stats();
            assert_eq!(stats.active_allocations, 1);
        }

        // Buffer should be returned to pool
        let stats = manager.get_stats();
        assert_eq!(stats.active_allocations, 0);
        assert_eq!(
            stats.pool_stats.f32_pool_size, 1,
            "dropping a managed buffer must hand its storage back to the pool"
        );
    }

    #[test]
    fn dropped_buffer_is_reused_on_the_next_request() {
        let manager = MemoryManager::new();
        drop(manager.get_f32_buffer(4096));

        let reused = manager.get_f32_buffer(4096);
        let stats = manager.get_stats();
        assert!(
            stats.pool_hit_rate > 0.0,
            "second request of the same size did not hit the pool"
        );
        assert_eq!(stats.pool_stats.f32_pool_size, 0);
        assert_eq!(stats.active_allocations, 1);
        drop(reused);
    }

    /// Regression: `into_inner` used to take the buffer out of the `Option`
    /// itself, so `Drop` found nothing to do and never decremented
    /// `active_allocations`. `MemoryStats` then reported a leak that grew by one
    /// for every buffer that was ever handed out.
    #[test]
    fn into_inner_balances_active_allocations() {
        let manager = MemoryManager::new();

        for _ in 0..3 {
            let buffer = manager.get_f32_buffer(1024);
            assert_eq!(manager.get_stats().active_allocations, 1);
            let raw = buffer.into_inner();
            assert!(raw.capacity() >= 1024);
            assert_eq!(
                manager.get_stats().active_allocations,
                0,
                "into_inner left a phantom allocation behind"
            );
        }
    }

    #[test]
    fn into_inner_keeps_the_storage_out_of_the_pool() {
        let manager = MemoryManager::new();
        let escaped = manager.get_f32_buffer(4096).into_inner();

        let stats = manager.get_stats();
        assert_eq!(
            stats.pool_stats.f32_pool_size, 0,
            "storage handed to the caller must not also sit in the pool"
        );
        assert_eq!(stats.active_allocations, 0);
        drop(escaped);
    }

    #[test]
    fn unmanaged_buffer_release_does_not_underflow_stats() {
        let manager = MemoryManager::with_config(MemoryConfig {
            enable_pooling: false,
            ..Default::default()
        });

        let buffer = manager.get_f32_buffer(256);
        assert_eq!(buffer.get().capacity(), 256);
        drop(buffer.into_inner());

        assert_eq!(manager.get_stats().active_allocations, 0);
    }

    #[test]
    fn test_buffer_reuse() {
        let manager = MemoryManager::new();

        // Create and drop a buffer
        {
            let _buffer = manager.get_f32_buffer(1000);
        }

        // Get another buffer of similar size
        {
            let _buffer = manager.get_f32_buffer(800);
            let stats = manager.get_stats();
            // Should have some pool hit rate if buffer was reused
            assert!(stats.pool_hit_rate >= 0.0);
        }
    }

    #[test]
    fn test_pre_allocation() {
        let manager = MemoryManager::new();
        manager.pre_allocate_common_sizes();

        let stats = manager.get_stats();
        assert!(stats.pool_stats.total_pool_memory > 0);
    }

    #[test]
    fn test_memory_config() {
        let config = MemoryConfig {
            enable_pooling: false,
            max_pool_size: 5,
            ..Default::default()
        };

        let manager = MemoryManager::with_config(config);
        assert!(!manager.config().enable_pooling);
        assert_eq!(manager.config().max_pool_size, 5);
    }

    #[test]
    fn test_buffer_acquisition_releases_pool_before_waiting_for_stats() {
        let manager = Arc::new(MemoryManager::new());
        let stats_guard = manager.stats.lock().unwrap();
        let transition = Arc::new(Barrier::new(2));
        let (reached_tx, reached_rx) = mpsc::sync_channel(0);
        manager.set_after_buffer_pool_access_hook(reached_tx, transition.clone());

        let (done_tx, done_rx) = mpsc::channel();
        let worker_manager = manager.clone();
        let worker = thread::spawn(move || {
            drop(worker_manager.get_f32_buffer(128));
            done_tx.send(()).unwrap();
        });

        reached_rx.recv_timeout(LOCK_TEST_TIMEOUT).unwrap();
        let pool_was_released = manager.buffer_pools.try_lock().is_ok();
        transition.wait();
        drop(stats_guard);

        done_rx.recv_timeout(LOCK_TEST_TIMEOUT).unwrap();
        worker.join().unwrap();
        assert!(
            pool_was_released,
            "buffer acquisition held the pool mutex while waiting for stats"
        );
    }

    #[test]
    fn test_stats_snapshot_releases_stats_before_waiting_for_pool() {
        let manager = Arc::new(MemoryManager::new());
        let pools_guard = manager.buffer_pools.lock().unwrap();
        let transition = Arc::new(Barrier::new(2));
        let (reached_tx, reached_rx) = mpsc::sync_channel(0);
        manager.set_after_stats_snapshot_hook(reached_tx, transition.clone());

        let (done_tx, done_rx) = mpsc::channel();
        let worker_manager = manager.clone();
        let worker = thread::spawn(move || {
            worker_manager.get_stats();
            done_tx.send(()).unwrap();
        });

        reached_rx.recv_timeout(LOCK_TEST_TIMEOUT).unwrap();
        let stats_were_released = manager.stats.try_lock().is_ok();
        transition.wait();
        drop(pools_guard);

        done_rx.recv_timeout(LOCK_TEST_TIMEOUT).unwrap();
        worker.join().unwrap();
        assert!(
            stats_were_released,
            "statistics snapshot held the stats mutex while waiting for pools"
        );
    }

    #[test]
    fn test_allocate_block_smoke() {
        let manager = MemoryManager::new();
        let block = manager.allocate_block(64, 8);
        assert!(block.is_ok());
        drop(block);
    }
}