//! GPU-resident trajectory tensor batch for PPO training.
//!
//! Replaces `trajectories::TrajectoryTensors` which used `candle_core::Tensor`.
//! All data stored as `CudaSlice<f32>` with zero Candle overhead.

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

use ml_core::MLError;
use ml_core::action_space::FactoredAction;

use super::GpuContext;
use super::tensor_util::{CudaVec, cuda_from_slice};

/// GPU-resident tensors for a trajectory batch.
///
/// All data is in F32 on GPU. Actions are stored as u32 indices.
#[derive(Debug)]
pub struct CudaTrajectoryTensors {
    /// States `[batch_size, state_dim]` flattened row-major
    pub states: CudaVec,
    /// Action indices `[batch_size]` as u32
    pub actions: CudaSlice<u32>,
    /// Old log probabilities `[batch_size]`
    pub log_probs: CudaVec,
    /// Value estimates `[batch_size]`
    pub values: CudaVec,
    /// Advantages `[batch_size]`
    pub advantages: CudaVec,
    /// Returns (value targets) `[batch_size]`
    pub returns: CudaVec,
    /// Batch size
    pub batch_size: usize,
    /// State dimension
    pub state_dim: usize,
}

impl CudaTrajectoryTensors {
    /// Upload trajectory batch data to GPU.
    ///
    /// Takes the pre-flattened host-side data from `TrajectoryBatch` and uploads
    /// it to GPU as contiguous `CudaSlice` buffers.
    pub fn from_batch(
        ctx: &GpuContext,
        states_flat: &[f32],
        actions: &[FactoredAction],
        log_probs: &[f32],
        values: &[f32],
        advantages: &[f32],
        returns: &[f32],
        state_dim: usize,
    ) -> Result<Self, MLError> {
        let batch_size = log_probs.len();
        let stream = &ctx.stream;

        // Upload states
        let states = cuda_from_slice(stream, states_flat)?;

        // Convert actions to u32 indices and upload
        let action_indices: Vec<u32> = actions.iter().map(|a| a.to_index() as u32).collect();
        let mut actions_gpu = stream.alloc_zeros::<u32>(batch_size).map_err(|e| {
            MLError::TensorOperationError(format!("alloc actions: {e}"))
        })?;
        stream.memcpy_htod(&action_indices, &mut actions_gpu).map_err(|e| {
            MLError::TensorOperationError(format!("upload actions: {e}"))
        })?;

        // Upload scalar arrays
        let log_probs_gpu = cuda_from_slice(stream, log_probs)?;
        let values_gpu = cuda_from_slice(stream, values)?;
        let advantages_gpu = cuda_from_slice(stream, advantages)?;
        let returns_gpu = cuda_from_slice(stream, returns)?;

        Ok(Self {
            states,
            actions: actions_gpu,
            log_probs: log_probs_gpu,
            values: values_gpu,
            advantages: advantages_gpu,
            returns: returns_gpu,
            batch_size,
            state_dim,
        })
    }

    /// Upload from a mini-batch (subset of full batch).
    pub fn from_mini_batch(
        ctx: &GpuContext,
        states: &[Vec<f32>],
        actions: &[FactoredAction],
        log_probs: &[f32],
        values: &[f32],
        advantages: &[f32],
        returns: &[f32],
        state_dim: usize,
    ) -> Result<Self, MLError> {
        // Flatten states
        let mut states_flat = Vec::with_capacity(states.len() * state_dim);
        for state in states {
            states_flat.extend_from_slice(state);
        }

        Self::from_batch(ctx, &states_flat, actions, log_probs, values, advantages, returns, state_dim)
    }

    /// Extract a GPU-resident sub-batch via `DtoD` copy. Zero CPU involvement.
    ///
    /// `start` and `end` are sample indices (0-based). States are sliced
    /// as `[start * state_dim .. end * state_dim]`. All other fields
    /// are sliced as `[start .. end]`.
    pub fn sub_batch(
        &self,
        start: usize,
        end: usize,
        stream: &std::sync::Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let count = end.saturating_sub(start);
        if count == 0 || end > self.batch_size {
            return Err(MLError::TensorOperationError(format!(
                "sub_batch range [{start}..{end}) out of bounds for batch_size {}",
                self.batch_size
            )));
        }

        let states = d2d_subrange(stream, &self.states.data, start * self.state_dim, count * self.state_dim)?;
        let actions = d2d_subrange_u32(stream, &self.actions, start, count)?;
        let log_probs = d2d_subrange(stream, &self.log_probs.data, start, count)?;
        let values = d2d_subrange(stream, &self.values.data, start, count)?;
        let advantages = d2d_subrange(stream, &self.advantages.data, start, count)?;
        let returns = d2d_subrange(stream, &self.returns.data, start, count)?;

        Ok(Self {
            states: CudaVec::new(states, count * self.state_dim),
            actions,
            log_probs: CudaVec::new(log_probs, count),
            values: CudaVec::new(values, count),
            advantages: CudaVec::new(advantages, count),
            returns: CudaVec::new(returns, count),
            batch_size: count,
            state_dim: self.state_dim,
        })
    }
}

/// GPU-to-GPU copy of `count` f32 elements starting at `offset` from `src`.
fn d2d_subrange(
    stream: &std::sync::Arc<CudaStream>,
    src: &CudaSlice<f32>,
    offset: usize,
    count: usize,
) -> Result<CudaSlice<f32>, MLError> {
    let mut dst = stream
        .alloc_zeros::<f32>(count)
        .map_err(|e| MLError::TensorOperationError(format!("d2d_subrange alloc({count}): {e}")))?;
    let src_view = src.slice(offset..offset + count);
    let nbytes = count * std::mem::size_of::<f32>();
    {
        use cudarc::driver::{DevicePtr, DevicePtrMut};
        let (src_ptr, _sg) = src_view.device_ptr(stream);
        let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
        // SAFETY: src and dst are valid GPU allocations of sufficient size.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, nbytes, stream.cu_stream(),
            )
            .map_err(|e| MLError::TensorOperationError(format!("d2d_subrange copy: {e}")))?;
        }
    }
    Ok(dst)
}

/// GPU-to-GPU copy of `count` u32 elements starting at `offset` from `src`.
fn d2d_subrange_u32(
    stream: &std::sync::Arc<CudaStream>,
    src: &CudaSlice<u32>,
    offset: usize,
    count: usize,
) -> Result<CudaSlice<u32>, MLError> {
    let mut dst = stream
        .alloc_zeros::<u32>(count)
        .map_err(|e| MLError::TensorOperationError(format!("d2d_subrange_u32 alloc({count}): {e}")))?;
    let src_view = src.slice(offset..offset + count);
    let nbytes = count * std::mem::size_of::<u32>();
    {
        use cudarc::driver::{DevicePtr, DevicePtrMut};
        let (src_ptr, _sg) = src_view.device_ptr(stream);
        let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
        // SAFETY: src and dst are valid GPU allocations of sufficient size.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr, src_ptr, nbytes, stream.cu_stream(),
            )
            .map_err(|e| MLError::TensorOperationError(format!("d2d_subrange_u32 copy: {e}")))?;
        }
    }
    Ok(dst)
}
