//! GPU tensor utilities — allocation helpers and `CudaVec` wrapper.

use std::sync::Arc;

use cudarc::driver::{CudaSlice, CudaStream};

use ml_core::MLError;

/// A GPU-resident vector with known length.
///
/// Wraps `CudaSlice<f32>` with an explicit element count so callers can
/// reshape/index without a device round-trip.
#[derive(Debug, Clone)]
pub struct CudaVec {
    pub data: CudaSlice<f32>,
    pub len: usize,
}

impl CudaVec {
    pub const fn new(data: CudaSlice<f32>, len: usize) -> Self {
        Self { data, len }
    }

    /// Download to host `Vec<f32>`.
    #[cold]
    pub fn to_vec(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
        let mut host = vec![0.0_f32; self.len];
        stream.memcpy_dtoh(&self.data, &mut host).map_err(|e| { // gpu-exit: API boundary readback
            MLError::TensorOperationError(format!("DtoH copy failed: {e}"))
        })?;
        Ok(host)
    }

    /// Convert to a `GpuTensor` with shape metadata.
    pub fn to_gpu_tensor(
        &self,
        stream: &Arc<CudaStream>,
        shape: Vec<usize>,
    ) -> Result<ml_core::cuda_autograd::GpuTensor, MLError> {
        let host = self.to_vec(stream)?;
        ml_core::cuda_autograd::GpuTensor::from_host(&host, shape, stream)
    }
}

/// Allocate a zero-initialized GPU buffer.
pub fn cuda_zeros(stream: &Arc<CudaStream>, len: usize) -> Result<CudaVec, MLError> {
    let data = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to alloc zeros({len}): {e}"))
    })?;
    Ok(CudaVec::new(data, len))
}

/// Allocate a GPU buffer filled with ones.
pub fn cuda_ones(stream: &Arc<CudaStream>, len: usize) -> Result<CudaVec, MLError> {
    let host = vec![1.0_f32; len];
    let mut data = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to alloc ones({len}): {e}"))
    })?;
    stream.memcpy_htod(&host, &mut data).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to upload ones: {e}"))
    })?;
    Ok(CudaVec::new(data, len))
}

/// Upload a host slice to GPU.
pub fn cuda_from_slice(stream: &Arc<CudaStream>, data: &[f32]) -> Result<CudaVec, MLError> {
    let len = data.len();
    let mut gpu = stream.alloc_zeros::<f32>(len).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to alloc from_slice({len}): {e}"))
    })?;
    stream.memcpy_htod(data, &mut gpu).map_err(|e| {
        MLError::TensorOperationError(format!("Failed to upload slice: {e}"))
    })?;
    Ok(CudaVec::new(gpu, len))
}
