#![allow(unsafe_code)] // Required for CUDA kernel launches and mapped-pinned alloc.

//! `GpuHealthDiag` — orchestrator for the HEALTH_DIAG GPU-port kernel family.
//!
//! Phase 2A landing: the orchestrator owns the `MappedHealthDiagSnapshot`
//! (a single mapped-pinned `HealthDiagSnapshot`) and the `CudaFunction`
//! handles for the kernel family that fills it. The Phase 2A commit holds
//! only the `health_diag_isv_mirror` kernel (single-block single-thread,
//! ISV signal-bus scalar copies). Subsequent phases append further kernel
//! handles to the same orchestrator (eval-histogram, q-mag reduction,
//! per-sample reduction, finalise — each tracked separately in
//! `docs/health_diag_inventory.md` under the Phase 2 commit plan).
//!
//! Phase 3 wires the launch chain into the captured graph alongside the
//! existing producer family (`launch_h_s2_rms_ema`, `launch_aux_heads_loss_ema`,
//! `launch_iqn_quantile_ema`, `launch_vsn_mask_ema`, etc.). Phase 4 deletes
//! the CPU-side reduction code and replaces the format-string emit with a
//! direct read of `snapshot.host_ref()`.
//!
//! Architecture rules upheld here (see `feedback_no_htod_htoh_only_mapped_pinned.md`,
//! `feedback_no_atomicadd.md`, `feedback_wire_everything_up.md`):
//!
//! * The snapshot lives in mapped-pinned memory — the only permitted
//!   CPU↔GPU communication path. No `memcpy_dtoh`, ever.
//! * Kernels write the snapshot through its device pointer; the host
//!   reads through `host_ref()` after stream-sync coherence.
//! * No `atomicAdd` in the kernel family (single-block, single-thread
//!   for `isv_mirror`; shmem tree reductions for the masking-heavy
//!   members in 2B-2D).
//! * The struct itself is built at construction (Phase 1, already
//!   merged); the wire-up to `GpuDqnTrainer` lands in this commit.
//!
//! Pearls applied:
//!
//! * `pearl_cold_path_no_exception_to_gpu_drives.md` — the snapshot is
//!   filled from device-resident producers, not from CPU-side
//!   re-derivation. Even ISV scalars (which the host can already read
//!   via `read_isv_signal_at`) are mirrored on the GPU side so the host
//!   emit path is allocator-free + DtoH-free.
//! * `feedback_no_partial_refactor.md` — the orchestrator is built per-
//!   commit but the kernel family converges before Phase 3's wire-up
//!   into the captured graph. CPU code stays in place during Phases
//!   2A-2E; Phase 4 deletes it in one shot once every snapshot field is
//!   populated.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

use crate::MLError;

use super::health_diag::MappedHealthDiagSnapshot;

// ── Cubin embedding ──────────────────────────────────────────────────────

/// Compiled cubin for `health_diag_kernel.cu`. Currently holds the Phase 2A
/// `health_diag_isv_mirror` kernel; phases 2B-2E append further kernel
/// entry points to the same `.cu` file (one cubin → many `load_function`
/// calls). Mirrors `AUX_HEADS_CUBIN` / `AUX_HEADS_LOSS_EMA_CUBIN`'s
/// embedding pattern at the top of `gpu_dqn_trainer.rs`.
static HEALTH_DIAG_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/health_diag_kernel.cubin"
));

/// Orchestrator owning the snapshot + kernel handles.
///
/// One instance per `GpuDqnTrainer`. Construction loads the cubin, resolves
/// every kernel symbol the family needs, and allocates the mapped-pinned
/// snapshot. Dropping releases all CUDA resources via cudarc + the
/// `MappedHealthDiagSnapshot::Drop` impl.
#[allow(missing_debug_implementations)]
pub(crate) struct GpuHealthDiag {
    /// Mapped-pinned `HealthDiagSnapshot`. Kernel writes via `dev_ptr()`;
    /// host emit reads via `host_ref()`.
    snapshot: MappedHealthDiagSnapshot,

    /// Phase 2A kernel handle. Single-block, single-thread; copies a
    /// curated set of ISV signal-bus slots into the corresponding
    /// snapshot fields. Args: ISV dev ptr, snapshot dev ptr, then 16
    /// `int` ISV indices (one per copy group).
    isv_mirror_kernel: CudaFunction,
}

impl GpuHealthDiag {
    /// Load the cubin, resolve every kernel symbol the family needs, and
    /// allocate the mapped-pinned snapshot. Mirrors
    /// `AuxHeadsForwardOps::new`'s loader shape.
    ///
    /// # Safety
    /// Caller must hold an active CUDA context on the current thread (the
    /// `MappedHealthDiagSnapshot::new` allocation requires it).
    pub(crate) unsafe fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();

        let module = context
            .load_cubin(HEALTH_DIAG_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("health_diag cubin load: {e}")))?;
        let isv_mirror_kernel = module
            .load_function("health_diag_isv_mirror")
            .map_err(|e| MLError::ModelError(format!("health_diag_isv_mirror load: {e}")))?;

        let snapshot = MappedHealthDiagSnapshot::new()
            .map_err(MLError::ModelError)?;

        Ok(Self {
            snapshot,
            isv_mirror_kernel,
        })
    }

    /// Borrow the snapshot for host-side reading. The caller must have
    /// stream-synchronised any prior kernel-family launch before treating
    /// the values as authoritative — mapped-pinned coherence guarantees
    /// visibility after the sync barrier.
    #[inline]
    pub(crate) fn snapshot(&self) -> &MappedHealthDiagSnapshot {
        &self.snapshot
    }

    /// Phase 2A — launch `health_diag_isv_mirror`.
    ///
    /// Copies the curated ISV signal-bus slots listed in the kernel doc-
    /// comment into the corresponding `HealthDiagSnapshot` fields. The
    /// caller passes the ISV indices (`REWARD_POPART_EMA_INDEX`,
    /// `MOE_EXPERT_UTIL_EMA_BASE`, etc.) directly so the kernel never
    /// embeds the numbering — single source of truth lives in
    /// `gpu_dqn_trainer.rs`.
    ///
    /// `isv_dev_ptr` MUST be the trainer's `isv_signals_dev_ptr` (the
    /// mapped-pinned ISV bus's device pointer). Non-zero invariant
    /// enforced by the caller — this method does not re-check it (loud
    /// failure preferred over silent skip; matches `launch_h_s2_rms_ema`
    /// debug-assert pattern).
    ///
    /// Cost: ~22 global loads + 22 stores from a single thread.
    /// Comfortably under any kernel-launch overhead.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_isv_mirror(
        &self,
        stream: &Arc<CudaStream>,
        isv_dev_ptr: u64,
        reward_popart_idx: i32,
        reward_cf_idx: i32,
        reward_trail_idx: i32,
        reward_micro_idx: i32,
        reward_opp_cost_idx: i32,
        reward_bonus_idx: i32,
        vsn_mag_idx: i32,
        vsn_dir_idx: i32,
        target_drift_mag_idx: i32,
        target_drift_dir_idx: i32,
        aux_next_bar_mse_idx: i32,
        aux_regime_ce_idx: i32,
        // SP13 B1.0 (2026-05-05): `aux_label_scale_idx: i32` parameter
        // dropped together with ISV[117] retirement.
        moe_expert_util_base_idx: i32,
        moe_gate_entropy_idx: i32,
        moe_lambda_eff_idx: i32,
    ) -> Result<(), MLError> {
        let snap_dev_ptr: u64 = self.snapshot.dev_ptr();
        unsafe {
            stream
                .launch_builder(&self.isv_mirror_kernel)
                .arg(&isv_dev_ptr)
                .arg(&snap_dev_ptr)
                .arg(&reward_popart_idx)
                .arg(&reward_cf_idx)
                .arg(&reward_trail_idx)
                .arg(&reward_micro_idx)
                .arg(&reward_opp_cost_idx)
                .arg(&reward_bonus_idx)
                .arg(&vsn_mag_idx)
                .arg(&vsn_dir_idx)
                .arg(&target_drift_mag_idx)
                .arg(&target_drift_dir_idx)
                .arg(&aux_next_bar_mse_idx)
                .arg(&aux_regime_ce_idx)
                // SP13 B1.0 (2026-05-05): aux_label_scale_idx arg dropped
                // — kernel signature shortened in lockstep.
                .arg(&moe_expert_util_base_idx)
                .arg(&moe_gate_entropy_idx)
                .arg(&moe_lambda_eff_idx)
                .launch(LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("health_diag_isv_mirror: {e}")))?;
        }
        Ok(())
    }
}
