use cudarc::driver::{CudaDevice, CudaSlice, DeviceSlice};

use super::{GpuError, gpu_unwrap};

/// Device-only buffer (faster, but not CPU-accessible)
pub struct DeviceBuffer<T: cudarc::driver::DeviceRepr> {
    inner: CudaSlice<T>,
}

impl<T: cudarc::driver::DeviceRepr + Clone> DeviceBuffer<T> {
    /// Allocate uninitialized
    pub fn new(len: usize) -> Result<Self, GpuError> {
        let ctx = gpu_unwrap();
        let inner = unsafe { ctx.device.alloc::<T>(len)? };
        Ok(Self { inner })
    }

    /// Allocate and copy from host slice
    pub fn from_slice(data: &[T]) -> Result<Self, GpuError> {
        let ctx = gpu_unwrap();
        let inner = ctx.device.htod_sync_copy(data)?;
        Ok(Self { inner })
    }

    /// Copy back to host
    pub fn to_vec(&self) -> Result<Vec<T>, GpuError> {
        let ctx = gpu_unwrap();
        Ok(ctx.device.dtoh_sync_copy(&self.inner)?)
    }

    /// Raw device pointer (for kernel params)
    pub fn as_ptr(&self) -> *mut T {
        *self.inner.device_ptr() as *mut T
    }

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

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

    /// Get the underlying CudaSlice (for cudarc APIs)
    pub fn as_cuda_slice(&self) -> &CudaSlice<T> {
        &self.inner
    }

    pub fn as_cuda_slice_mut(&mut self) -> &mut CudaSlice<T> {
        &mut self.inner
    }
}

/// Unified memory buffer (CPU + GPU accessible)
/// Note: cudarc doesn't have built-in unified memory,
/// so we use raw CUDA calls or just use DeviceBuffer + explicit copies
pub type UnifiedBuffer<T> = DeviceBuffer<T>;
