//! Staging buffer pool management.

use memkit_gpu::{MkGpuBackend, MkBufferUsage, MkMemoryType};
use parking_lot::Mutex;
use std::collections::VecDeque;

/// A pool of reusable staging buffers.
pub struct StagingPool<B: MkGpuBackend> {
    buffer_size: usize,
    max_buffers: usize,
    free_buffers: Mutex<VecDeque<B::BufferHandle>>,
    total_created: Mutex<usize>,
}

impl<B: MkGpuBackend> StagingPool<B> {
    /// Create a new staging pool.
    pub fn new(buffer_size: usize, max_buffers: usize) -> Self {
        Self {
            buffer_size,
            max_buffers,
            free_buffers: Mutex::new(VecDeque::with_capacity(max_buffers)),
            total_created: Mutex::new(0),
        }
    }
    
    /// Get the size of buffers in this pool.
    pub fn buffer_size(&self) -> usize {
        self.buffer_size
    }
    
    /// Get the number of free buffers available.
    pub fn free_count(&self) -> usize {
        self.free_buffers.lock().len()
    }
    
    /// Get the total number of buffers created.
    pub fn total_created(&self) -> usize {
        *self.total_created.lock()
    }
    
    /// Acquire a staging buffer from the pool, or create a new one.
    pub fn acquire(&self, backend: &B) -> Result<StagingHandle<B>, B::Error> {
        // Try to get from pool first
        if let Some(handle) = self.free_buffers.lock().pop_front() {
            return Ok(StagingHandle {
                handle,
                size: self.buffer_size,
                from_pool: true,
            });
        }
        
        // Create new buffer
        let handle = backend.create_buffer(
            self.buffer_size,
            MkBufferUsage::TRANSFER_SRC,
            MkMemoryType::HostVisible,
        )?;
        
        *self.total_created.lock() += 1;
        
        Ok(StagingHandle {
            handle,
            size: self.buffer_size,
            from_pool: true,
        })
    }
    
    /// Release a staging buffer back to the pool.
    pub fn release(&self, staging: StagingHandle<B>, backend: &B) {
        let mut free = self.free_buffers.lock();
        
        if free.len() < self.max_buffers {
            free.push_back(staging.handle);
        } else {
            // Pool is full, destroy the buffer
            backend.destroy_buffer(&staging.handle);
        }
    }
    
    /// Clear all pooled buffers.
    pub fn clear(&self, backend: &B) {
        let mut free = self.free_buffers.lock();
        for handle in free.drain(..) {
            backend.destroy_buffer(&handle);
        }
    }
    
    /// Pre-populate the pool with buffers.
    pub fn warm(&self, count: usize, backend: &B) -> Result<(), B::Error> {
        let to_create = count.min(self.max_buffers);
        let mut free = self.free_buffers.lock();
        
        for _ in 0..to_create {
            if free.len() >= self.max_buffers {
                break;
            }
            
            let handle = backend.create_buffer(
                self.buffer_size,
                MkBufferUsage::TRANSFER_SRC,
                MkMemoryType::HostVisible,
            )?;
            
            *self.total_created.lock() += 1;
            free.push_back(handle);
        }
        
        Ok(())
    }
}

/// Handle to a staging buffer acquired from the pool.
pub struct StagingHandle<B: MkGpuBackend> {
    pub(crate) handle: B::BufferHandle,
    pub(crate) size: usize,
    pub(crate) from_pool: bool,
}

impl<B: MkGpuBackend> StagingHandle<B> {
    /// Get the raw buffer handle.
    pub fn handle(&self) -> &B::BufferHandle {
        &self.handle
    }
    
    /// Get the buffer size.
    pub fn size(&self) -> usize {
        self.size
    }
    
    /// Write data to the staging buffer.
    pub fn write_data<T: Copy>(&self, backend: &B, data: &[T]) -> bool {
        let byte_size = std::mem::size_of_val(data);
        if byte_size > self.size {
            return false;
        }
        
        if let Some(ptr) = backend.map(&self.handle) {
            unsafe {
                std::ptr::copy_nonoverlapping(
                    data.as_ptr() as *const u8,
                    ptr,
                    byte_size,
                );
            }
            backend.unmap(&self.handle);
            true
        } else {
            false
        }
    }
    
    /// Write raw bytes to the staging buffer.
    pub fn write_bytes(&self, backend: &B, data: &[u8]) -> bool {
        if data.len() > self.size {
            return false;
        }
        
        if let Some(ptr) = backend.map(&self.handle) {
            unsafe {
                std::ptr::copy_nonoverlapping(data.as_ptr(), ptr, data.len());
            }
            backend.unmap(&self.handle);
            true
        } else {
            false
        }
    }
}

/// A dynamically-sized staging pool that can provide buffers of various sizes.
pub struct DynamicStagingPool<B: MkGpuBackend> {
    /// Small buffers (up to 16KB)
    small_pool: StagingPool<B>,
    /// Medium buffers (up to 64KB)
    medium_pool: StagingPool<B>,
    /// Large buffers (up to 256KB)
    large_pool: StagingPool<B>,
}

impl<B: MkGpuBackend> DynamicStagingPool<B> {
    /// Create a new dynamic staging pool with default sizes.
    pub fn new() -> Self {
        Self {
            small_pool: StagingPool::new(16 * 1024, 32),      // 16KB x 32
            medium_pool: StagingPool::new(64 * 1024, 16),     // 64KB x 16
            large_pool: StagingPool::new(256 * 1024, 8),      // 256KB x 8
        }
    }
    
    /// Acquire a staging buffer that can hold at least `size` bytes.
    pub fn acquire(&self, size: usize, backend: &B) -> Result<StagingHandle<B>, B::Error> {
        if size <= 16 * 1024 {
            self.small_pool.acquire(backend)
        } else if size <= 64 * 1024 {
            self.medium_pool.acquire(backend)
        } else if size <= 256 * 1024 {
            self.large_pool.acquire(backend)
        } else {
            // Too large for pool, create one-off buffer
            let handle = backend.create_buffer(
                size,
                MkBufferUsage::TRANSFER_SRC,
                MkMemoryType::HostVisible,
            )?;
            
            Ok(StagingHandle {
                handle,
                size,
                from_pool: false,
            })
        }
    }
    
    /// Release a staging buffer back to the appropriate pool.
    pub fn release(&self, staging: StagingHandle<B>, backend: &B) {
        if !staging.from_pool {
            // Not from pool, just destroy it
            backend.destroy_buffer(&staging.handle);
            return;
        }
        
        if staging.size <= 16 * 1024 {
            self.small_pool.release(staging, backend);
        } else if staging.size <= 64 * 1024 {
            self.medium_pool.release(staging, backend);
        } else if staging.size <= 256 * 1024 {
            self.large_pool.release(staging, backend);
        } else {
            backend.destroy_buffer(&staging.handle);
        }
    }
    
    /// Clear all pools.
    pub fn clear(&self, backend: &B) {
        self.small_pool.clear(backend);
        self.medium_pool.clear(backend);
        self.large_pool.clear(backend);
    }
}

impl<B: MkGpuBackend> Default for DynamicStagingPool<B> {
    fn default() -> Self {
        Self::new()
    }
}
