//! Per-worker buffer pool with small and large queues (lock-free).
//!
//! Design goals:
//! - minimize allocations by reusing `BytesMut`
//! - small/large split to avoid wasting large buffers on tiny requests
//! - safe for async usage via lock-free queues
//!
//! Notes:
//! - `buffer_size` is the small buffer size (e.g. 4K/8K)
//! - `max_buffer_size` caps large buffers (default 1_048_576)
//! - buffers are recycled at end of request/connection

use bytes::BytesMut;
use crossbeam_queue::SegQueue;
use std::sync::Arc;

const DEFAULT_MAX_BUFFER_SIZE: usize = 1_048_576;

/// Per-worker buffer pool.
#[derive(Clone, Debug)]
pub struct BufferPool {
    small: Arc<SegQueue<BytesMut>>,
    large: Arc<SegQueue<BytesMut>>,
    buffer_size: usize,
    max_buffer_size: usize,
}

impl BufferPool {
    /// Create a new pool.
    pub fn new(buffer_size: usize, max_buffer_size: usize) -> Self {
        let max_buffer_size = max_buffer_size.max(buffer_size);
        Self {
            small: Arc::new(SegQueue::new()),
            large: Arc::new(SegQueue::new()),
            buffer_size,
            max_buffer_size,
        }
    }

    /// Create a pool with default max buffer size (1_048_576).
    pub fn with_default_max(buffer_size: usize) -> Self {
        Self::new(buffer_size, DEFAULT_MAX_BUFFER_SIZE)
    }

    /// Acquire a buffer sized for at least `min_size`.
    ///
    /// If `min_size <= buffer_size`, a small buffer is used.
    /// Otherwise, a large buffer is used (capped at `max_buffer_size`).
    pub fn acquire(&self, min_size: usize) -> BytesMut {
        let target = if min_size <= self.buffer_size {
            self.buffer_size
        } else {
            min_size.min(self.max_buffer_size).max(self.buffer_size)
        };

        if target <= self.buffer_size {
            if let Some(mut buf) = self.small.pop() {
                buf.clear();
                if buf.capacity() < target {
                    buf.reserve(target - buf.capacity());
                }
                return buf;
            }
        } else if let Some(mut buf) = self.large.pop() {
            buf.clear();
            if buf.capacity() < target {
                buf.reserve(target - buf.capacity());
            }
            return buf;
        }

        BytesMut::with_capacity(target)
    }

    /// Release a buffer back into the pool.
    pub fn release(&self, mut buf: BytesMut) {
        buf.clear();

        let cap = buf.capacity();
        if cap <= self.buffer_size {
            self.small.push(buf);
        } else if cap <= self.max_buffer_size {
            self.large.push(buf);
        } else {
            // Drop oversized buffers to avoid pool bloat.
        }
    }

    /// Buffer size used for the small pool.
    pub fn buffer_size(&self) -> usize {
        self.buffer_size
    }

    /// Maximum buffer size for the large pool.
    pub fn max_buffer_size(&self) -> usize {
        self.max_buffer_size
    }
}

/// RAII wrapper that returns the buffer to its pool on drop.
///
/// IMPORTANT: For reading into a `BytesMut`, you generally want to write into its
/// *spare capacity* and then advance the length. This wrapper provides helpers
/// for that pattern.
#[derive(Debug)]
pub struct PooledBuffer {
    pool: BufferPool,
    buffer: Option<BytesMut>,
}

impl PooledBuffer {
    /// Create a new pooled buffer from the given pool.
    pub fn new(pool: BufferPool, min_size: usize) -> Self {
        let buffer = pool.acquire(min_size);
        Self {
            pool,
            buffer: Some(buffer),
        }
    }

    /// Ensure the inner buffer has at least `additional` spare capacity.
    pub fn reserve_spare(&mut self, additional: usize) {
        let buf = self
            .buffer
            .as_mut()
            .expect("buffer already taken from PooledBuffer");
        buf.reserve(additional);
    }

    /// Get a mutable slice of the buffer's spare capacity to read into.
    ///
    /// This returns a `&mut [u8]` pointing to the uninitialized tail region.
    /// After writing `n` bytes into the returned slice, call `advance_len(n)`
    /// to make those bytes part of the buffer.
    pub fn spare_capacity_mut(&mut self) -> &mut [u8] {
        let buf = self
            .buffer
            .as_mut()
            .expect("buffer already taken from PooledBuffer");
        let len = buf.len();
        let cap = buf.capacity();

        if cap == len {
            // Grow at least by 1 to expose spare capacity.
            buf.reserve(1);
        }

        // SAFETY:
        // - `BytesMut` uses a contiguous allocation.
        // - We expose the region [len..cap) as a mutable slice for IO reads.
        // - The bytes are considered uninitialized until `advance_len` is called.
        unsafe {
            let ptr = buf.as_mut_ptr().add(len);
            let spare = buf.capacity() - len;
            std::slice::from_raw_parts_mut(ptr, spare)
        }
    }

    /// Advance the buffer length after writing `n` bytes into the spare capacity.
    pub fn advance_len(&mut self, n: usize) {
        let buf = self
            .buffer
            .as_mut()
            .expect("buffer already taken from PooledBuffer");
        let len = buf.len();
        let cap = buf.capacity();
        let new_len = len + n;
        assert!(
            new_len <= cap,
            "advance_len would exceed capacity (len={}, n={}, cap={})",
            len,
            n,
            cap
        );

        // SAFETY: we just wrote `n` bytes into the spare capacity.
        unsafe {
            buf.set_len(new_len);
        }
    }

    /// Access the inner buffer mutably (initialized region only).
    ///
    /// Prefer `spare_capacity_mut` + `advance_len` for IO reads.
    pub fn buffer_mut(&mut self) -> &mut BytesMut {
        self.buffer
            .as_mut()
            .expect("buffer already taken from PooledBuffer")
    }

    /// Access the inner buffer immutably (initialized region only).
    pub fn buffer(&self) -> &BytesMut {
        self.buffer
            .as_ref()
            .expect("buffer already taken from PooledBuffer")
    }

    /// Take ownership of the buffer without returning it to the pool.
    pub fn take(mut self) -> BytesMut {
        self.buffer
            .take()
            .expect("buffer already taken from PooledBuffer")
    }
}

impl Drop for PooledBuffer {
    fn drop(&mut self) {
        if let Some(buf) = self.buffer.take() {
            self.pool.release(buf);
        }
    }
}

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

    #[test]
    fn test_pool_acquire_release_small_buffer() {
        let pool = BufferPool::new(1024, 4096);
        let mut buf = pool.acquire(100);
        assert!(buf.capacity() >= 1024);
        buf.extend_from_slice(b"payload");
        pool.release(buf);

        let reused = pool.acquire(64);
        assert!(reused.capacity() >= 1024);
        assert_eq!(reused.len(), 0);
    }

    #[test]
    fn test_pool_acquire_large_buffer() {
        let pool = BufferPool::new(1024, 4096);
        let buf = pool.acquire(3000);
        assert!(buf.capacity() >= 3000);
    }

    #[test]
    fn test_pool_max_buffer_size_is_clamped() {
        let pool = BufferPool::new(4096, 1024);
        assert_eq!(pool.max_buffer_size(), 4096);
    }

    #[test]
    fn test_pooled_buffer_spare_and_advance() {
        let pool = BufferPool::new(256, 2048);
        let mut pooled = PooledBuffer::new(pool, 128);
        let dst = pooled.spare_capacity_mut();
        dst[..4].copy_from_slice(b"test");
        pooled.advance_len(4);
        assert_eq!(&pooled.buffer()[..], b"test");
    }
}
