//! Raw CUDA neural network primitives replacing `candle_nn`.
//!
//! All forward passes use cuBLAS sgemm for matrix multiplication and
//! custom CUDA kernels for activations (`ReLU`, softmax, tanh).
//! Weights are stored as `CudaSlice<f32>` with no Candle overhead.
//! Adam optimizer runs entirely on GPU via element-wise CUDA kernels.

#![allow(unsafe_code)] // Required for CUDA kernel launches and cuBLAS FFI

pub mod linear;
pub mod activations;
pub mod adam;
pub mod lstm;
pub mod networks;
pub mod softmax;
pub mod tensor_util;
pub mod trajectory_tensors;

pub use linear::CudaLinear;
pub use activations::{cuda_relu, cuda_relu_backward, cuda_tanh, cuda_sigmoid};
pub use adam::CudaAdam;
pub use lstm::CudaLSTM;
pub use networks::{CudaPolicyNetwork, CudaValueNetwork};
pub use softmax::{cuda_softmax, cuda_log_softmax};
pub use tensor_util::{CudaVec, cuda_zeros, cuda_ones, cuda_from_slice};
pub use trajectory_tensors::CudaTrajectoryTensors;

use std::mem::ManuallyDrop;
use std::sync::Arc;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
use ml_core::MLError;

/// Shared CUDA context for all neural network operations.
///
/// Wraps a cuBLAS handle and CUDA stream for efficient GPU computation.
/// All layers share the same `GpuContext` to avoid redundant allocations.
#[derive(Clone)]
pub struct GpuContext {
    pub stream: Arc<CudaStream>,
    pub cublas: Arc<cudarc::cublas::CudaBlas>,
}

impl GpuContext {
    /// Create a new GPU context on the default device.
    pub fn new() -> Result<Self, MLError> {
        let cuda_ctx = cudarc::driver::CudaContext::new(0).map_err(|e| {
            MLError::InitializationError {
                component: "GpuContext".to_owned(),
                message: format!("Failed to create CUDA context: {e}"),
            }
        })?;
        let stream = cuda_ctx.default_stream();
        let cublas = Arc::new(
            cudarc::cublas::CudaBlas::new(stream.clone()).map_err(|e| {
                MLError::InitializationError {
                    component: "GpuContext".to_owned(),
                    message: format!("Failed to create cuBLAS handle: {e}"),
                }
            })?,
        );
        Ok(Self { stream, cublas })
    }

    /// Create from an existing stream (for interop with existing CUDA pipeline).
    pub fn from_stream(stream: Arc<CudaStream>) -> Result<Self, MLError> {
        let cublas = Arc::new(
            cudarc::cublas::CudaBlas::new(stream.clone()).map_err(|e| {
                MLError::InitializationError {
                    component: "GpuContext".to_owned(),
                    message: format!("Failed to create cuBLAS handle: {e}"),
                }
            })?,
        );
        Ok(Self { stream, cublas })
    }
}

impl std::fmt::Debug for GpuContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GpuContext").finish()
    }
}

/// Extract raw CUDA device pointer from a `CudaSlice`.
///
/// Uses `ManuallyDrop` on the `SyncOnDrop` guard to skip read event recording.
/// Safe when operating on the same stream (CUDA ordering guaranteed).
pub(crate) fn raw_ptr(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

/// Extract raw mutable CUDA device pointer from a `CudaSlice`.
pub(crate) fn raw_ptr_mut(slice: &mut CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr_mut(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}
