pub mod dtype;
pub mod shape;
pub mod random;

use crate::hardware::buffer::Buffer;
use crate::hardware::device::DeviceType;
pub use crate::tensor::dtype::DType;
pub use crate::tensor::shape::Shape;

pub struct Tensor {
    buffer: Buffer,
    shape: Shape,
    dtype: DType,
}

impl Tensor {
    pub fn new(shape: Shape, dtype: DType, device: DeviceType) -> Result<Self, String> {
        let size_bytes = shape.num_elements() * dtype.size_of();
        let buffer = Buffer::new(device, size_bytes)?;
        
        Ok(Tensor {
            buffer,
            shape,
            dtype,
        })
    }

    pub fn shape(&self) -> &Shape {
        &self.shape
    }

    pub fn dtype(&self) -> DType {
        self.dtype
    }

    pub fn device(&self) -> DeviceType {
        self.buffer.device_type()
    }
    
    pub fn fill_random(&mut self, seed: u64) {
        if self.device() != DeviceType::Cpu {
            // Cannot fill random directly on non-CPU yet
            return;
        }

        use crate::tensor::random::XorShift64;
        let mut rng = XorShift64::new(seed);

        match self.dtype {
            DType::F32 => {
                let count = self.shape.num_elements();
                let mut data = Vec::with_capacity(count);
                for _ in 0..count {
                    data.push(rng.next_f32());
                }
                unsafe {
                    let _ = self.buffer.write(&data);
                }
            }
            _ => {
                // TODO: Implement for other types
            }
        }
    }

    pub fn write_data<T: Copy>(&mut self, data: &[T]) -> Result<(), String> {
        // TODO: check if T size matches dtype size
        unsafe {
            self.buffer.write(data)
        }
    }

    pub fn read_data<T: Copy>(&self) -> Result<Vec<T>, String> {
        // TODO: check if T size matches dtype size
        let count = self.shape.num_elements();
        unsafe {
            self.buffer.read(count)
        }
    }
}