#![allow(unsafe_code)]

//! GPU-resident training guard and Q-value monitor.
//!
//! Wraps four CUDA kernels from `training_guard_kernel.cu` to perform
//! NaN/Inf detection, loss clipping, gradient collapse checks, and
//! Q-value statistics entirely on GPU.
//!
//! Per-step results are read via **mapped pinned memory** with
//! double-buffering. The kernel writes to one buffer while the CPU
//! reads the previous step's result from the other — no `memcpy_dtoh`
//! is needed on the hot path.
//!
//! The accumulator buffer (3 floats) stays in device memory; a single
//! `memcpy_dtoh` at epoch boundary returns (mean_loss, mean_grad_norm).

use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use std::ffi::c_void;
use std::mem::MaybeUninit;
use std::sync::{Arc, OnceLock};

use ml_core::nvtx::NvtxRange;
use crate::MLError;

// ── PTX cache ──────────────────────────────────────────────────────────────

static TRAINING_GUARD_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

fn compile_training_guard_ptx(context: &CudaContext) -> Result<Ptx, String> {
    let kernel_src = include_str!("training_guard_kernel.cu");
    crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context)
}

// ── Mapped pinned memory helper ───────────────────────────────────────────

/// CPU+GPU visible buffer allocated via `cuMemHostAlloc(DEVICEMAP)`.
/// The kernel writes via `dev_ptr`, the CPU reads via `host_ptr`.
/// No `memcpy_dtoh` is needed -- `__threadfence_system()` in the kernel
/// ensures write visibility across PCIe.
struct MappedBuffer {
    host_ptr: *mut f32,
    dev_ptr: cudarc::driver::sys::CUdeviceptr,
    len: usize,
}

// Safety: MappedBuffer is only accessed from the thread that created
// the GpuTrainingGuard, which holds a CUDA context. The host_ptr is
// valid for the lifetime of the allocation.
unsafe impl Send for MappedBuffer {}
unsafe impl Sync for MappedBuffer {}

impl MappedBuffer {
    /// Allocate `len` floats of mapped pinned memory.
    ///
    /// # Safety
    /// Caller must ensure a CUDA context is active on the current thread.
    unsafe fn new(len: usize) -> Result<Self, MLError> {
        let num_bytes = len * std::mem::size_of::<f32>();
        let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
            | cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;

        let host_ptr = cudarc::driver::result::malloc_host(num_bytes, flags)
            .map_err(|e| MLError::ModelError(format!("mapped host alloc ({len} f32): {e}")))?
            as *mut f32;

        // Zero-init so first read returns safe defaults
        std::ptr::write_bytes(host_ptr, 0, len);

        // Get device-visible pointer for this mapped host allocation
        let mut dev_ptr_raw = MaybeUninit::uninit();
        cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
            dev_ptr_raw.as_mut_ptr(),
            host_ptr as *mut c_void,
            0,
        )
        .result()
        .map_err(|e| MLError::ModelError(format!("cuMemHostGetDevicePointer: {e}")))?;

        Ok(Self {
            host_ptr,
            dev_ptr: dev_ptr_raw.assume_init(),
            len,
        })
    }

    /// Read the buffer contents from the CPU side.
    ///
    /// Uses `read_volatile` to prevent the CPU from caching stale values.
    /// Caller must ensure that the kernel writing to this buffer has completed
    /// (e.g., via double-buffering with sufficient pipeline delay).
    ///
    /// # Panics (debug builds)
    /// Panics if `idx >= self.len`. All call sites use compile-time-known
    /// indices that are within bounds (verified by construction).
    fn read(&self, idx: usize) -> f32 {
        debug_assert!(
            idx < self.len,
            "MappedBuffer::read OOB: idx={idx} >= len={}",
            self.len
        );
        // Safety: host_ptr is valid for self.len floats; all callers use
        // hardcoded indices < len (guard=0..5 of 7, qstats=0..3 of 4,
        // qdiv=0..4 of 5). debug_assert catches bugs in dev builds.
        unsafe { std::ptr::read_volatile(self.host_ptr.add(idx)) }
    }
}

impl Drop for MappedBuffer {
    fn drop(&mut self) {
        // Safety: host_ptr was returned by malloc_host and is non-null.
        unsafe {
            // Intentionally ignoring the result -- nothing useful to do on
            // free failure in a destructor.
            #[allow(clippy::let_underscore_must_use)]
            let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void);
        }
    }
}

// ── Result types ───────────────────────────────────────────────────────────

/// Result from `training_guard_check`: per-step safety flags and scalar values.
#[derive(Debug, Clone)]
pub struct GuardResult {
    /// True if loss or grad_norm is NaN or Inf.
    pub halt_nan: bool,
    /// True if loss exceeded the clip threshold.
    pub halt_loss_clip: bool,
    /// True if grad_norm fell below the collapse threshold (and not in warmup).
    pub halt_grad_collapse: bool,
    /// Loss after clipping (min(loss, clip_threshold)); NaN/Inf pass through.
    pub clipped_loss: f32,
    /// Raw loss as read from the GPU scalar.
    pub raw_loss: f32,
    /// Raw gradient norm as read from the GPU scalar.
    pub raw_grad_norm: f32,
}

/// Per-batch Q-value statistics from `qvalue_stats_reduce`.
#[derive(Debug, Clone)]
pub struct QValueStats {
    /// Min of per-sample max-Q across the batch.
    pub q_min: f32,
    /// Max of per-sample max-Q across the batch.
    pub q_max: f32,
    /// Mean of per-sample max-Q across the batch.
    pub q_mean: f32,
    /// Mean of all Q-values (all actions x all samples; collapse detection).
    pub q_all_mean: f32,
}

/// Single-sample Q-value divergence result from `qvalue_divergence_check`.
#[derive(Debug, Clone)]
pub struct QValueDivergence {
    /// Minimum Q-value.
    pub q_min: f32,
    /// Maximum Q-value.
    pub q_max: f32,
    /// Mean Q-value.
    pub q_mean: f32,
    /// Variance of Q-values.
    pub q_variance: f32,
    /// True if |q_min| or |q_max| exceeds the divergence threshold.
    pub divergence_detected: bool,
}

// ── Main struct ────────────────────────────────────────────────────────────

/// GPU-resident training guard + Q-value monitor.
///
/// Uses mapped pinned memory with double-buffering to eliminate
/// `memcpy_dtoh` from the per-step hot path. The kernel writes to
/// one buffer while the CPU reads the other (one-step delay).
pub struct GpuTrainingGuard {
    fused_check_accum_func: CudaFunction,
    qvalue_stats_func: CudaFunction,
    qvalue_div_func: CudaFunction,

    /// Double-buffered mapped output for `training_guard_check` (7 floats each).
    guard_mapped: [MappedBuffer; 2],
    /// Current write buffer index (alternates 0/1).
    guard_buf_idx: usize,
    /// Whether at least one step has completed (first step has no previous data).
    guard_has_prev: bool,

    /// Double-buffered mapped output for `qvalue_stats_reduce` (4 floats each).
    qstats_mapped: [MappedBuffer; 2],
    qstats_buf_idx: usize,
    qstats_has_prev: bool,

    /// Double-buffered mapped output for `qvalue_divergence_check` (5 floats each).
    qdiv_mapped: [MappedBuffer; 2],
    qdiv_buf_idx: usize,
    qdiv_has_prev: bool,

    /// Accumulator buffer (3 floats: loss_sum, grad_norm_sum, step_count).
    /// Stays in device memory -- kernel does read-modify-write.
    /// Single `memcpy_dtoh` at boundary only (not per-step).
    acc_buf: CudaSlice<f32>,

    /// Welford running-mean accumulator for Q-value estimation (CPU-side, no GPU tensor).
    q_count: usize,
    q_mean_val: f64,

    stream: Arc<CudaStream>,
}

impl GpuTrainingGuard {
    /// Create a new `GpuTrainingGuard` from a CUDA stream.
    ///
    /// Compiles the PTX once per process (OnceLock cached), loads all four
    /// kernel functions, and pre-allocates all mapped double-buffers.
    pub fn new(stream: Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();

        // Compile PTX (once per process via OnceLock, disk-cached via nvcc)
        let ptx_result = TRAINING_GUARD_PTX.get_or_init(|| compile_training_guard_ptx(&context));
        let ptx = ptx_result.as_ref().map_err(|e| {
            MLError::ModelError(format!("training_guard PTX: {e}"))
        })?;
        let module = context.load_module(ptx.clone()).map_err(|e| {
            MLError::ModelError(format!("training_guard module load: {e}"))
        })?;

        let fused_check_accum_func = module
            .load_function("training_guard_check_and_accumulate")
            .map_err(|e| MLError::ModelError(format!("training_guard_check_and_accumulate load: {e}")))?;
        let qvalue_stats_func = module
            .load_function("qvalue_stats_reduce")
            .map_err(|e| MLError::ModelError(format!("qvalue_stats_reduce load: {e}")))?;
        let qvalue_div_func = module
            .load_function("qvalue_divergence_check")
            .map_err(|e| MLError::ModelError(format!("qvalue_divergence_check load: {e}")))?;

        // Allocate device accumulator buffer (stays on-device, read at epoch boundary)
        let acc_buf = stream.alloc_zeros::<f32>(3).map_err(|e| {
            MLError::ModelError(format!("alloc acc_buf: {e}"))
        })?;

        // Allocate mapped pinned double-buffers
        // Safety: CUDA context is active (we just loaded the module on this device).
        let guard_mapped = unsafe { [MappedBuffer::new(7)?, MappedBuffer::new(7)?] };
        let qstats_mapped = unsafe { [MappedBuffer::new(4)?, MappedBuffer::new(4)?] };
        let qdiv_mapped = unsafe { [MappedBuffer::new(5)?, MappedBuffer::new(5)?] };

        Ok(Self {
            fused_check_accum_func,
            qvalue_stats_func,
            qvalue_div_func,
            guard_mapped,
            guard_buf_idx: 0,
            guard_has_prev: false,
            qstats_mapped,
            qstats_buf_idx: 0,
            qstats_has_prev: false,
            qdiv_mapped,
            qdiv_buf_idx: 0,
            qdiv_has_prev: false,
            acc_buf,
            q_count: 0,
            q_mean_val: 0.0,
            stream,
        })
    }

    /// Constructor from a CUDA stream (primary API).
    ///
    /// Alias for `new()` — kept for call-site compatibility during migration.
    pub fn from_stream(stream: Arc<CudaStream>) -> Result<Self, MLError> {
        Self::new(stream)
    }

    /// Run the guard check and accumulate kernels for one training step.
    ///
    /// `loss_gpu` and `grad_norm_gpu` must be F32 `CudaSlice<f32>` scalars (1 element each).
    ///
    /// Returns the safety flags and scalar values from the *previous* step
    /// (one-step delay due to double-buffering). On the very first call,
    /// returns safe defaults (no halts, zero loss/grad_norm).
    pub fn check_and_accumulate(
        &mut self,
        loss_gpu: &CudaSlice<f32>,
        grad_norm_gpu: &CudaSlice<f32>,
        clip_threshold: f32,
        collapse_threshold: f32,
        warmup: bool,
    ) -> Result<GuardResult, MLError> {
        let _nvtx = NvtxRange::new("gpu_training_guard");

        // Read PREVIOUS step's result (one-step delay; first step returns defaults)
        let prev_result = if self.guard_has_prev {
            let read_idx = 1 - self.guard_buf_idx;
            let g = &self.guard_mapped[read_idx];
            GuardResult {
                halt_nan: g.read(0) != 0.0,
                halt_loss_clip: g.read(1) != 0.0,
                halt_grad_collapse: g.read(2) != 0.0,
                clipped_loss: g.read(3),
                raw_loss: g.read(4),
                raw_grad_norm: g.read(5),
            }
        } else {
            // First step: no previous data, return safe defaults
            GuardResult {
                halt_nan: false,
                halt_loss_clip: false,
                halt_grad_collapse: false,
                clipped_loss: 0.0,
                raw_loss: 0.0,
                raw_grad_norm: 0.0,
            }
        };

        // Get current write buffer's device pointer (kernel writes here)
        let write_dev_ptr = self.guard_mapped[self.guard_buf_idx].dev_ptr;

        let warmup_int: i32 = if warmup { 1 } else { 0 };

        // Fused kernel: training_guard_check_and_accumulate -- grid=(1,1,1), block=(1,1,1)
        let single_cfg = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (1, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: loss_gpu and grad_norm_gpu are valid F32 CudaSlice on the same context.
        unsafe {
            self.stream
                .launch_builder(&self.fused_check_accum_func)
                .arg(loss_gpu)
                .arg(grad_norm_gpu)
                .arg(&write_dev_ptr)
                .arg(&mut self.acc_buf)
                .arg(&clip_threshold)
                .arg(&collapse_threshold)
                .arg(&warmup_int)
                .launch(single_cfg)
                .map_err(|e| MLError::ModelError(format!("training_guard_check_and_accumulate launch: {e}")))?;
        }

        // Swap buffer index for next step
        self.guard_has_prev = true;
        self.guard_buf_idx = 1 - self.guard_buf_idx;

        Ok(prev_result)
    }

    /// Read the epoch-boundary loss and grad_norm averages from the accumulator.
    pub fn read_accumulators(&mut self) -> Result<(f64, f64), MLError> {
        let mut host = [0.0_f32; 3];
        self.stream.memcpy_dtoh(&self.acc_buf, &mut host).map_err(|e| { // gpu-exit: 12B epoch accumulators
            MLError::ModelError(format!("acc_buf readback: {e}"))
        })?;

        let loss_sum = host[0] as f64;
        let grad_sum = host[1] as f64;
        let steps = host[2] as f64;

        if steps <= 0.0 {
            return Ok((0.0, 0.0));
        }

        Ok((loss_sum / steps, grad_sum / steps))
    }

    /// Zero the accumulator buffer on-device (call at epoch start/end).
    pub fn reset_accumulators(&mut self) -> Result<(), MLError> {
        self.stream.memset_zeros(&mut self.acc_buf).map_err(|e| {
            MLError::ModelError(format!("reset acc_buf: {e}"))
        })?;
        Ok(())
    }

    /// Compute Q-value statistics over a full batch via `qvalue_stats_reduce`.
    ///
    /// `q_values` must be a contiguous F32 `CudaSlice<f32>` with `batch_size * num_actions` elements.
    pub fn qvalue_stats(
        &mut self,
        q_values: &CudaSlice<f32>,
        batch_size: usize,
        num_actions: usize,
    ) -> Result<QValueStats, MLError> {
        if batch_size == 0 || num_actions == 0 {
            return Ok(QValueStats { q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_all_mean: 0.0 });
        }

        let prev_stats = if self.qstats_has_prev {
            let read_idx = 1 - self.qstats_buf_idx;
            let g = &self.qstats_mapped[read_idx];
            QValueStats { q_min: g.read(0), q_max: g.read(1), q_mean: g.read(2), q_all_mean: g.read(3) }
        } else {
            QValueStats { q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_all_mean: 0.0 }
        };

        let write_dev_ptr = self.qstats_mapped[self.qstats_buf_idx].dev_ptr;
        let bs_i32 = batch_size as i32;
        let na_i32 = num_actions as i32;

        let cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };

        // Safety: q_values has batch_size * num_actions elements.
        unsafe {
            self.stream
                .launch_builder(&self.qvalue_stats_func)
                .arg(q_values)
                .arg(&write_dev_ptr)
                .arg(&bs_i32)
                .arg(&na_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("qvalue_stats_reduce launch: {e}")))?;
        }

        self.qstats_has_prev = true;
        self.qstats_buf_idx = 1 - self.qstats_buf_idx;

        Ok(prev_stats)
    }

    /// Check a single sample for Q-value divergence via `qvalue_divergence_check`.
    ///
    /// `q_values` must be a contiguous F32 `CudaSlice<f32>` with `num_actions` elements.
    pub fn qvalue_divergence(
        &mut self,
        q_values: &CudaSlice<f32>,
        num_actions: usize,
        threshold: f32,
    ) -> Result<QValueDivergence, MLError> {
        if num_actions == 0 {
            return Ok(QValueDivergence {
                q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_variance: 0.0, divergence_detected: false,
            });
        }

        let prev_div = if self.qdiv_has_prev {
            let read_idx = 1 - self.qdiv_buf_idx;
            let g = &self.qdiv_mapped[read_idx];
            QValueDivergence {
                q_min: g.read(0), q_max: g.read(1), q_mean: g.read(2),
                q_variance: g.read(3), divergence_detected: g.read(4) != 0.0,
            }
        } else {
            QValueDivergence {
                q_min: 0.0, q_max: 0.0, q_mean: 0.0, q_variance: 0.0, divergence_detected: false,
            }
        };

        let write_dev_ptr = self.qdiv_mapped[self.qdiv_buf_idx].dev_ptr;
        let na_i32 = num_actions as i32;
        let single_cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };

        // Safety: q_values has num_actions elements.
        unsafe {
            self.stream
                .launch_builder(&self.qvalue_div_func)
                .arg(q_values)
                .arg(&write_dev_ptr)
                .arg(&na_i32)
                .arg(&threshold)
                .launch(single_cfg)
                .map_err(|e| MLError::ModelError(format!("qvalue_divergence_check launch: {e}")))?;
        }

        self.qdiv_has_prev = true;
        self.qdiv_buf_idx = 1 - self.qdiv_buf_idx;

        Ok(prev_div)
    }

    /// Accumulate a Q-value mean using Welford running mean (CPU-side arithmetic).
    ///
    /// The caller reads the scalar Q-value from GPU once (already done for logging)
    /// and passes it here. No GPU tensor operations needed for the accumulator itself.
    pub fn accumulate_q_value(&mut self, avg_q: f32) {
        self.q_count += 1;
        let delta = avg_q as f64 - self.q_mean_val;
        self.q_mean_val += delta / self.q_count as f64;
    }

    /// Read accumulated Q-value mean at epoch end.
    pub fn read_q_accumulator(&self) -> f64 {
        if self.q_count == 0 {
            return 0.0;
        }
        self.q_mean_val
    }

    /// Reset Q-value accumulator for new epoch.
    pub fn reset_q_accumulator(&mut self) {
        self.q_count = 0;
        self.q_mean_val = 0.0;
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ptx_compilation() {
        let context = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
        let stream = context.default_stream();
        let context = stream.context();
        let result = compile_training_guard_ptx(&context);
        if let Err(ref e) = result {
            panic!("training_guard PTX compilation failed: {e}");
        }
    }
}
