use crate::error::{self, Result};
use crate::ffi;
use crate::pool;

/// RAII wrapper around a GPU memory allocation. Dropping returns the buffer to
/// the per-device caching pool (see [`crate::pool`]) for reuse instead of
/// calling `cudaFree`.
///
/// `len` is the logical (requested) length used for all copy/memset bounds and
/// exposed via [`GpuBuffer::len`]; `cap` is the physical size class the pool
/// rounded up to (>= `len`), used only to bucket the buffer for reuse. The
/// extra `cap - len` bytes are never exposed to callers, so pooling is
/// numerically transparent. `device` records which device pool to return to.
pub struct GpuBuffer {
    ptr: *mut u8,
    len: usize,
    cap: usize,
    device: i32,
}

impl GpuBuffer {
    /// Allocate at least `len` bytes on the calling thread's current device,
    /// reusing a pooled buffer when one of the matching size class is free.
    /// The contents are **uninitialized** (a reused buffer holds stale bytes);
    /// callers that need zeros must memset (see [`crate::Storage::zeros`]).
    pub fn alloc(len: usize) -> Result<Self> {
        assert!(len > 0, "cannot allocate 0 bytes on GPU");
        let a = pool::acquire(len)?;
        Ok(Self {
            ptr: a.ptr,
            len,
            cap: a.cap,
            device: a.device,
        })
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn as_ptr(&self) -> *const u8 {
        self.ptr
    }

    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.ptr
    }

    /// Copy data from a host (CPU) slice to this GPU buffer.
    pub fn copy_from_host(&mut self, src: &[u8]) -> Result<()> {
        assert!(src.len() <= self.len, "source larger than buffer");
        error::check(unsafe {
            ffi::cudaMemcpy(self.ptr, src.as_ptr(), src.len(), ffi::CUDA_MEMCPY_H2D)
        })
    }

    /// Copy data from this GPU buffer to a host (CPU) slice.
    pub fn copy_to_host(&self, dst: &mut [u8]) -> Result<()> {
        assert!(dst.len() <= self.len, "destination larger than buffer");
        error::check(unsafe {
            ffi::cudaMemcpy(dst.as_mut_ptr(), self.ptr, dst.len(), ffi::CUDA_MEMCPY_D2H)
        })
    }

    /// Set every byte of the buffer to `value` on the device (no host copy).
    /// Used to zero op-output buffers without a blocking H2D memcpy of zeros.
    pub fn memset(&mut self, value: u8) -> Result<()> {
        error::check(unsafe { ffi::cudaMemset(self.ptr, value as i32, self.len) })
    }
}

impl Drop for GpuBuffer {
    fn drop(&mut self) {
        // Return to the device pool for reuse (no cudaFree). The pool retains
        // the raw pointer for the process lifetime; on process exit the OS
        // reclaims the device context, so this is not a leak.
        pool::release(self.ptr, self.device, self.cap);
    }
}

unsafe impl Send for GpuBuffer {}