#![allow(unsafe_code)]

//! Per-stream cuBLAS + cublasLt handle registry for the DQN training pipeline.
//!
//! # NVIDIA cuBLAS §2.1.4 — Results Reproducibility
//!
//! > "This guarantee holds when a single CUDA stream is active only. If multiple
//! > concurrent streams are active, the library may optimize total performance
//! > by picking different internal implementations."
//!
//! Foxhunt's fused DQN training forks onto `iqn_stream`, `attn_stream`, etc.
//! The NVIDIA-documented remediation is: **one cuBLAS handle per stream**.
//!
//! This module implements that remediation. [`PerStreamCublasHandles`] owns a
//! registry of `(cublasLtHandle_t, workspace)` pairs keyed by raw CUDA stream
//! pointer. Each cublasLtMatmul launch dispatches on the handle+workspace bound
//! to *its own* stream — never sharing a handle across concurrent streams.
//!
//! ## Composition with Option B (DeterministicAlgoSelector)
//!
//! Option B (commit `bb399b635`) pins deterministic algo selection via
//! `AlgoGetIds`+`AlgoInit`+`AlgoCheck`. That pinning is per-call (it runs
//! against whichever lt_handle is passed in), so it composes cleanly with the
//! per-stream registry: the same shape key produces the same algo ID on any
//! lt_handle on the same hardware.
//!
//! ## Workspace isolation
//!
//! Each per-stream handle owns its own 32 MB `CudaSlice<u8>` workspace. Sharing
//! a workspace across concurrent streams races; isolating them removes that
//! variance.

use std::collections::HashMap;
use std::sync::Arc;

use cudarc::cublas::result as cublas_result;
use cudarc::cublas::sys as cublas_sys;
use cudarc::cublaslt::result as cublaslt_result;
use cudarc::cublaslt::sys as cublaslt_sys;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream};
use parking_lot::RwLock;

use crate::MLError;

// ── Handle wrappers ───────────────────────────────────────────────────────────

/// Wrapper to make `cublasHandle_t` Send + Sync.
///
/// Used by `gpu_curiosity_trainer` (which owns a dedicated one-stream-one-handle
/// pair outside this registry) and by [`PerStreamCublasHandles`] itself for
/// classic-cuBLAS GEMMs that the cuBLAS-Lt heuristic refuses to dispatch (e.g.
/// TLOB tiny attention shapes M=16, K=32). `Drop` destroys the handle via
/// `cublasDestroy`.
pub(crate) struct SendSyncCublasHandle(pub(crate) cublas_sys::cublasHandle_t);

unsafe impl Send for SendSyncCublasHandle {}
unsafe impl Sync for SendSyncCublasHandle {}

impl Drop for SendSyncCublasHandle {
    fn drop(&mut self) {
        unsafe {
            let _ = cublas_result::destroy_handle(self.0);
        }
    }
}

/// Wrapper to make `cublasLtHandle_t` Send + Sync.
///
/// Owned via `Arc` but only ever accessed from the single GPU trainer thread.
/// `Drop` destroys the underlying cublasLt handle via `cublasLtDestroy`.
pub(crate) struct SendSyncCublasLtHandle(pub(crate) cublaslt_sys::cublasLtHandle_t);

unsafe impl Send for SendSyncCublasLtHandle {}
unsafe impl Sync for SendSyncCublasLtHandle {}

impl Drop for SendSyncCublasLtHandle {
    fn drop(&mut self) {
        unsafe {
            let _ = cublaslt_result::destroy_handle(self.0);
        }
    }
}

// ── Per-stream registry entry ─────────────────────────────────────────────────

/// One `(cublasLt handle, classic cublas handle, workspace)` triple bound to a
/// specific CUDA stream.
///
/// Stored inside the registry's `HashMap<StreamKey, PerStreamLt>`. The fields
/// are `pub(crate)` so other `cuda_pipeline` modules can read the raw pointers
/// without going through an accessor method (hot-path FFI).
///
/// Both handles share the same workspace (`cublasSetWorkspace_v2` + the
/// cublasLt workspace ptr arg accept identical device buffers); each handle is
/// already bound to this entry's stream at creation time and rebound on demand
/// via [`PerStreamCublasHandles::classic_for`] / [`PerStreamCublasHandles::lt_for`].
pub(crate) struct PerStreamLt {
    /// Raw cublasLt handle created for this specific stream. Dropped (via
    /// `cublasLtDestroy`) when the entry leaves the registry — which only
    /// happens when [`PerStreamCublasHandles`] itself is dropped.
    pub(crate) lt_handle: SendSyncCublasLtHandle,
    /// Raw classic cuBLAS handle created for this specific stream. Used for
    /// GEMMs whose dimensions are too narrow for cuBLAS-Lt's tensor-core
    /// heuristic (e.g. TLOB attention M=16, K=32). Dropped via `cublasDestroy`.
    pub(crate) classic_handle: SendSyncCublasHandle,
    /// Owned workspace buffer (kept alive for the device pointer's lifetime).
    #[allow(dead_code)]
    pub(crate) workspace_buf: CudaSlice<u8>,
    /// Raw device pointer for `cublasLtMatmul`'s `workspace` arg
    /// (also bound on `classic_handle` via `cublasSetWorkspace_v2`).
    pub(crate) workspace_ptr: u64,
    /// Workspace size in bytes.
    pub(crate) workspace_size: usize,
}

/// Per-stream cublasLt workspace size. H100 TF32 HMMA algorithms need this
/// much headroom for the heuristic's largest workspace.
const LT_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;

// ── PerStreamCublasHandles ────────────────────────────────────────────────────

/// Per-stream cublasLt handle registry.
///
/// Owns one cublasLt handle + workspace per CUDA stream that launches matmul.
/// Default-stream lt_handle/workspace are exposed as direct fields for the
/// common case — callers on the default stream pay zero hashmap cost.
#[allow(missing_debug_implementations)]
pub struct PerStreamCublasHandles {
    // ── Default-stream cublasLt (fast path) ──────────────────────────────────
    /// cublasLt handle bound to the default/main training stream. Exposed as a
    /// direct field so the overwhelmingly common default-stream call site pays
    /// no lookup cost.
    pub(crate) lt_handle: SendSyncCublasLtHandle,
    /// Classic cuBLAS handle bound to the default training stream.
    ///
    /// Created alongside [`lt_handle`]; shares the same workspace buffer.
    /// Used for GEMMs whose dimensions are too narrow for cuBLAS-Lt's
    /// tensor-core-oriented heuristic (TLOB attention is the canonical case
    /// — M=TLOB_OUT=16, K∈{16,32}, well below the cuBLAS-Lt sweet spot).
    /// This is the NVIDIA-recommended split: cuBLAS-Lt for large/
    /// tensor-core-optimised GEMMs, classic cuBLAS for general-purpose any-
    /// shape GEMMs.
    pub(crate) classic_handle: SendSyncCublasHandle,
    /// Owned workspace backing [`lt_workspace_ptr`].
    #[allow(dead_code)]
    pub(crate) _lt_workspace_buf: CudaSlice<u8>,
    /// Raw device pointer for the default-stream cublasLt workspace.
    pub(crate) lt_workspace_ptr: u64,
    /// Size of the default-stream cublasLt workspace in bytes.
    pub(crate) lt_workspace_size: usize,

    // ── Side-stream registry ────────────────────────────────────────────────
    /// Per-stream (handle, workspace) pairs for non-default streams. Keyed by
    /// raw CUstream pointer (`usize`). Populated lazily on first access; all
    /// entries are dropped together when `self` is dropped.
    registry: RwLock<HashMap<usize, PerStreamLt>>,

    // ── Bias-add kernels (stream-agnostic) ──────────────────────────────────
    /// `add_bias_relu_f32_kernel(output, bias, out_dim, total)` — f32 bias + ReLU.
    pub(crate) add_bias_relu_f32_kernel: CudaFunction,
    /// `add_bias_f32_kernel(output, bias, out_dim, total)` — f32 bias, no activation.
    pub(crate) add_bias_f32_kernel: CudaFunction,

    /// The default CUDA stream for this registry (main training stream).
    pub(crate) stream: Arc<CudaStream>,
}

impl PerStreamCublasHandles {
    /// Create a new registry with the default stream's cublasLt + classic
    /// cuBLAS handles, workspace, and bias-add kernels pre-provisioned.
    ///
    /// Side-stream handles are created lazily by [`Self::lt_for`] /
    /// [`Self::classic_for`].
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        // ── Default-stream cublasLt + classic handles + workspace ───────────
        let (lt_raw_handle, classic_raw_handle, lt_workspace_buf, lt_workspace_ptr) =
            create_handles_and_workspace(stream)?;

        // ── Bias-add kernels ─────────────────────────────────────────────────
        let (add_bias_relu_f32_kernel, add_bias_f32_kernel) = load_bias_kernels(stream)?;

        Ok(Self {
            lt_handle: SendSyncCublasLtHandle(lt_raw_handle),
            classic_handle: SendSyncCublasHandle(classic_raw_handle),
            _lt_workspace_buf: lt_workspace_buf,
            lt_workspace_ptr,
            lt_workspace_size: LT_WORKSPACE_BYTES,
            registry: RwLock::new(HashMap::new()),
            add_bias_relu_f32_kernel,
            add_bias_f32_kernel,
            stream: Arc::clone(stream),
        })
    }

    /// Return `(lt_handle, workspace_ptr, workspace_size)` for `stream`.
    ///
    /// - Fast path: if `stream` is the default stream, returns the cached
    ///   default-stream triple with no hashmap lookup.
    /// - Hot path: RwLock read for existing entries (O(1) hashmap).
    /// - Cold path: upgrades to write lock only on first access for a new
    ///   stream. First-access cost: one `cublasLtCreate` (~ms) + one 32 MB
    ///   device allocation. Subsequent calls are a single hashmap read.
    pub fn lt_for(
        &self,
        stream: &CudaStream,
    ) -> Result<(cublaslt_sys::cublasLtHandle_t, u64, usize), MLError> {
        let key = stream.cu_stream() as usize;

        // Fast path: default stream is the overwhelmingly common case.
        if key == self.stream.cu_stream() as usize {
            return Ok((
                self.lt_handle.0,
                self.lt_workspace_ptr,
                self.lt_workspace_size,
            ));
        }

        // Hot path: side stream already registered.
        {
            let reg = self.registry.read();
            if let Some(entry) = reg.get(&key) {
                return Ok((
                    entry.lt_handle.0,
                    entry.workspace_ptr,
                    entry.workspace_size,
                ));
            }
        }

        // Cold path: first access for this stream — provision a new handle.
        let mut reg = self.registry.write();
        // Re-check after acquiring the write lock (another caller may have
        // created it while we were upgrading).
        if let Some(entry) = reg.get(&key) {
            return Ok((
                entry.lt_handle.0,
                entry.workspace_ptr,
                entry.workspace_size,
            ));
        }

        // Allocate workspace via the default stream (same context as `stream`).
        // The workspace lives in device memory; once allocated, it's usable from
        // any stream on the same context. We only need an `Arc<CudaStream>` for
        // cudarc's `alloc_zeros`, and `self.stream` is always available.
        let (lt_handle, classic_handle, buf, ptr) =
            create_handles_and_workspace(&self.stream)?;
        let entry = PerStreamLt {
            lt_handle: SendSyncCublasLtHandle(lt_handle),
            classic_handle: SendSyncCublasHandle(classic_handle),
            workspace_buf: buf,
            workspace_ptr: ptr,
            workspace_size: LT_WORKSPACE_BYTES,
        };
        let result = (entry.lt_handle.0, entry.workspace_ptr, entry.workspace_size);
        reg.insert(key, entry);
        Ok(result)
    }

    /// Convenience: just the lt_handle for `stream` (for call sites that
    /// already know the workspace).
    pub fn lt_handle_for(
        &self,
        stream: &CudaStream,
    ) -> Result<cublaslt_sys::cublasLtHandle_t, MLError> {
        Ok(self.lt_for(stream)?.0)
    }

    /// Return the classic cuBLAS handle for `stream`, already bound to the
    /// stream and to the per-stream workspace via `cublasSetStream` +
    /// `cublasSetWorkspace_v2`. Caller can immediately invoke
    /// `cublasSgemm_v2` etc.
    ///
    /// Same fast-path / hot-path / cold-path layout as [`Self::lt_for`].
    /// Cold path (first call for a side stream) creates a fresh `cublasLt`
    /// handle alongside the classic one — they always come as a pair so
    /// later `lt_for` calls hit the hot path with zero extra cost.
    pub fn classic_for(
        &self,
        stream: &CudaStream,
    ) -> Result<(cublas_sys::cublasHandle_t, u64, usize), MLError> {
        let key = stream.cu_stream() as usize;

        // Fast path: default stream.
        if key == self.stream.cu_stream() as usize {
            bind_classic_handle(
                self.classic_handle.0,
                stream,
                self.lt_workspace_ptr,
                self.lt_workspace_size,
            )?;
            return Ok((
                self.classic_handle.0,
                self.lt_workspace_ptr,
                self.lt_workspace_size,
            ));
        }

        // Hot path: side stream already registered.
        {
            let reg = self.registry.read();
            if let Some(entry) = reg.get(&key) {
                bind_classic_handle(
                    entry.classic_handle.0,
                    stream,
                    entry.workspace_ptr,
                    entry.workspace_size,
                )?;
                return Ok((
                    entry.classic_handle.0,
                    entry.workspace_ptr,
                    entry.workspace_size,
                ));
            }
        }

        // Cold path: provision a fresh (lt, classic) pair for this stream.
        let mut reg = self.registry.write();
        if let Some(entry) = reg.get(&key) {
            bind_classic_handle(
                entry.classic_handle.0,
                stream,
                entry.workspace_ptr,
                entry.workspace_size,
            )?;
            return Ok((
                entry.classic_handle.0,
                entry.workspace_ptr,
                entry.workspace_size,
            ));
        }

        let (lt_handle, classic_handle, buf, ptr) =
            create_handles_and_workspace(&self.stream)?;
        let entry = PerStreamLt {
            lt_handle: SendSyncCublasLtHandle(lt_handle),
            classic_handle: SendSyncCublasHandle(classic_handle),
            workspace_buf: buf,
            workspace_ptr: ptr,
            workspace_size: LT_WORKSPACE_BYTES,
        };
        let result = (
            entry.classic_handle.0,
            entry.workspace_ptr,
            entry.workspace_size,
        );
        // Bind to the caller's stream (handle was created during
        // `create_handles_and_workspace` already bound to `self.stream`).
        bind_classic_handle(entry.classic_handle.0, stream, entry.workspace_ptr, entry.workspace_size)?;
        reg.insert(key, entry);
        Ok(result)
    }

    /// Convenience: just the classic handle for `stream`.
    pub fn classic_handle_for(
        &self,
        stream: &CudaStream,
    ) -> Result<cublas_sys::cublasHandle_t, MLError> {
        Ok(self.classic_for(stream)?.0)
    }

    /// Eagerly provision a per-stream handle + workspace for `stream`.
    ///
    /// Must be called before CUDA Graph capture begins — the first-access
    /// cold path inside [`Self::lt_for`] calls `cudaMalloc` (via cudarc's
    /// `alloc_zeros`), which fails with `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED`
    /// when invoked during capture. Pre-registering side streams at trainer
    /// construction time (before `training_begin`) ensures every subsequent
    /// `lt_for` call hits the hot path.
    ///
    /// No-op if `stream` is the default stream (already provisioned in `new`).
    /// No-op if `stream` is already registered.
    pub fn pre_register_stream(&self, stream: &Arc<CudaStream>) -> Result<(), MLError> {
        let key = stream.cu_stream() as usize;
        if key == self.stream.cu_stream() as usize {
            return Ok(());
        }
        {
            let reg = self.registry.read();
            if reg.contains_key(&key) {
                return Ok(());
            }
        }
        let mut reg = self.registry.write();
        if reg.contains_key(&key) {
            return Ok(());
        }
        // Provision via `self.stream` so we can use cudarc's Arc-requiring alloc.
        let (lt_handle, classic_handle, buf, ptr) =
            create_handles_and_workspace(&self.stream)?;
        reg.insert(
            key,
            PerStreamLt {
                lt_handle: SendSyncCublasLtHandle(lt_handle),
                classic_handle: SendSyncCublasHandle(classic_handle),
                workspace_buf: buf,
                workspace_ptr: ptr,
                workspace_size: LT_WORKSPACE_BYTES,
            },
        );
        Ok(())
    }
}

// ── Provisioning helpers ──────────────────────────────────────────────────────

/// Create a `(cublasLt, classic cuBLAS)` handle pair sharing a 32 MB workspace.
///
/// The workspace is allocated via `stream` (cudarc's `alloc_zeros` requires
/// `&Arc<CudaStream>`) but is usable from any stream on the same CUDA context
/// — which is why the cold path in [`PerStreamCublasHandles::lt_for`] /
/// [`PerStreamCublasHandles::classic_for`] allocates workspaces via
/// `self.stream` regardless of which side stream requested the lookup.
///
/// Both handles are bound to `stream` and the workspace at creation time:
/// - `cublasLt` workspace is passed per-call (already idiomatic).
/// - Classic cuBLAS handle: `cublasSetStream` + `cublasSetWorkspace_v2`.
///
/// Classic handles default to TF32 tensor-core math
/// (`CUBLAS_TF32_TENSOR_OP_MATH`) — same convention as
/// `gpu_curiosity_trainer::CuriosityGemm` and as the cublasLt path uses
/// (`CUBLAS_COMPUTE_32F_FAST_TF32`). Determinism on shared trunk paths is
/// already enforced upstream by the deterministic algo selector; the TLOB
/// classic-sgemm path is auxiliary attention head-tail with the same
/// run-to-run tolerance the rest of the pipeline accepts.
fn create_handles_and_workspace(
    stream: &Arc<CudaStream>,
) -> Result<
    (
        cublaslt_sys::cublasLtHandle_t,
        cublas_sys::cublasHandle_t,
        CudaSlice<u8>,
        u64,
    ),
    MLError,
> {
    let lt_handle = cublaslt_result::create_handle()
        .map_err(|e| MLError::ModelError(format!("cublasLtCreate: {e:?}")))?;

    let classic_handle = cublas_result::create_handle()
        .map_err(|e| MLError::ModelError(format!("cublasCreate (classic): {e:?}")))?;

    let workspace_buf = stream
        .alloc_zeros::<u8>(LT_WORKSPACE_BYTES)
        .map_err(|e| MLError::ModelError(format!("cublasLt workspace alloc: {e}")))?;
    let workspace_ptr = workspace_buf.raw_ptr();

    // Bind classic handle to its creation-time stream + workspace + TF32.
    bind_classic_handle(classic_handle, stream, workspace_ptr, LT_WORKSPACE_BYTES)?;

    unsafe {
        let status = cublas_sys::cublasSetMathMode(
            classic_handle,
            cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
        );
        if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            return Err(MLError::ModelError(format!(
                "cublasSetMathMode (classic): {status:?}"
            )));
        }
    }

    Ok((lt_handle, classic_handle, workspace_buf, workspace_ptr))
}

/// Bind a classic cuBLAS handle to `stream` + `workspace_ptr`. Idempotent;
/// safe to call before every batch of GEMMs to make the binding explicit.
fn bind_classic_handle(
    handle: cublas_sys::cublasHandle_t,
    stream: &CudaStream,
    workspace_ptr: u64,
    workspace_size: usize,
) -> Result<(), MLError> {
    unsafe {
        let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st;
        cublas_result::set_stream(handle, cu_stream)
            .map_err(|e| MLError::ModelError(format!("cublasSetStream (classic): {e:?}")))?;
        let status = cublas_sys::cublasSetWorkspace_v2(
            handle,
            workspace_ptr as *mut std::ffi::c_void,
            workspace_size,
        );
        if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            return Err(MLError::ModelError(format!(
                "cublasSetWorkspace_v2 (classic): {status:?}"
            )));
        }
    }
    Ok(())
}

// ── Kernel loader ─────────────────────────────────────────────────────────────

/// Precompiled bias kernels cubin (produced by build.rs — zero runtime nvcc).
static BIAS_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/bias_kernels.cubin"));

/// Load f32 bias kernels from the precompiled cubin.
fn load_bias_kernels(
    stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction), MLError> {
    let context = stream.context();
    let module = context
        .load_cubin(BIAS_CUBIN.to_vec())
        .map_err(|e| MLError::ModelError(format!("bias cubin load: {e}")))?;

    let add_bias_relu_f32 = module
        .load_function("add_bias_relu_f32_kernel")
        .map_err(|e| MLError::ModelError(format!("add_bias_relu_f32_kernel load: {e}")))?;
    let add_bias_f32 = module
        .load_function("add_bias_f32_kernel")
        .map_err(|e| MLError::ModelError(format!("add_bias_f32_kernel load: {e}")))?;

    Ok((add_bias_relu_f32, add_bias_f32))
}
