use super::device::DeviceType;
use super::memory::MemoryManager;
use std::slice;

pub struct Buffer {
    ptr: *mut u8,
    size: usize,
    device_type: DeviceType,
}

impl Buffer {
    pub fn new(device_type: DeviceType, size: usize) -> Result<Self, String> {
        let ptr = MemoryManager::allocate(device_type, size)?;
        Ok(Buffer {
            ptr,
            size,
            device_type,
        })
    }

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

    pub fn device_type(&self) -> DeviceType {
        self.device_type
    }

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

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

    // Unsafe because the caller must ensure T matches the underlying data
    pub unsafe fn write<T: Copy>(&mut self, data: &[T]) -> Result<(), String> {
        if self.device_type != DeviceType::Cpu {
            return Err("Direct write only supported on CPU for now".to_string());
        }
        let bytes_needed = data.len() * std::mem::size_of::<T>();
        if bytes_needed > self.size {
            return Err("Buffer too small".to_string());
        }
        
        std::ptr::copy_nonoverlapping(
            data.as_ptr() as *const u8, 
            self.ptr, 
            bytes_needed
        );
        Ok(())
    }

    pub unsafe fn read<T: Copy>(&self, count: usize) -> Result<Vec<T>, String> {
         if self.device_type != DeviceType::Cpu {
            return Err("Direct read only supported on CPU for now".to_string());
        }
        let bytes_needed = count * std::mem::size_of::<T>();
        if bytes_needed > self.size {
            return Err("Read size exceeds buffer".to_string());
        }

        let mut result = Vec::with_capacity(count);
        let src = slice::from_raw_parts(self.ptr as *const T, count);
        result.extend_from_slice(src);
        Ok(result)
    }
}

impl Drop for Buffer {
    fn drop(&mut self) {
        MemoryManager::deallocate(self.device_type, self.ptr, self.size);
    }
}