//! Zero-overhead CUDA kernel launch bypassing cudarc's `PushKernelArg`
//! trait dispatch and event tracking. Used in the step hot path where
//! cudarc's ~200us-per-launch overhead dominates at ~70 launches/step.
//!
//! cudarc stays for device init, module loading, and memory allocation.
//! Only the launch path is bypassed here.

use cudarc::driver::sys::{self, CUevent, CUfunction, CUresult, CUstream, CUgraphExec};
use std::ffi::c_void;

/// Maximum kernel arguments supported by a single launch.
/// 32 covers every kernel in the step pipeline (largest is ~24 args).
const MAX_ARGS: usize = 32;

/// Fixed-capacity kernel argument buffer.
///
/// All values are stored as `u64` in a flat array. For device pointers
/// (`CUdeviceptr = u64`) this is a direct store. For smaller scalars
/// (`i32`, `f32`), the value occupies the low bytes of the `u64` slot;
/// CUDA reads `sizeof(param)` bytes from the pointer we give it, which
/// on little-endian reads exactly the right bytes.
///
/// The `arg_ptrs` array is **not** stored inline because the struct
/// would become self-referential and invalidate on move. Instead, call
/// [`RawArgs::build_arg_ptrs`] to produce a stack-local pointer array
/// right before launch.
pub struct RawArgs {
    storage: [u64; MAX_ARGS],
    len: usize,
}

impl RawArgs {
    #[inline(always)]
    pub fn new() -> Self {
        Self {
            storage: [0u64; MAX_ARGS],
            len: 0,
        }
    }

    /// Reset the argument list for reuse on the next kernel.
    #[inline(always)]
    pub fn reset(&mut self) {
        self.len = 0;
    }

    /// Number of arguments currently stored.
    #[inline(always)]
    pub fn len(&self) -> usize {
        self.len
    }

    /// Whether the argument list is empty.
    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Push a device pointer argument (`CUdeviceptr = u64`).
    #[inline(always)]
    pub fn push_ptr(&mut self, ptr: u64) {
        debug_assert!(self.len < MAX_ARGS, "RawArgs overflow");
        self.storage[self.len] = ptr;
        self.len += 1;
    }

    /// Push a raw `u64` argument (device pointer or 64-bit scalar).
    #[inline(always)]
    pub fn push_u64(&mut self, val: u64) {
        debug_assert!(self.len < MAX_ARGS, "RawArgs overflow");
        self.storage[self.len] = val;
        self.len += 1;
    }

    /// Push an `i32` scalar argument.
    ///
    /// Stored in the low 4 bytes of a `u64` slot. CUDA reads
    /// `sizeof(int)` = 4 bytes from the pointer, which on
    /// little-endian is exactly the `i32` value.
    #[inline(always)]
    pub fn push_i32(&mut self, val: i32) {
        debug_assert!(self.len < MAX_ARGS, "RawArgs overflow");
        // Zero-extend to u64 so the upper bytes are clean.
        self.storage[self.len] = val as u32 as u64;
        self.len += 1;
    }

    /// Push an `f32` scalar argument.
    ///
    /// Stored as the bit pattern in the low 4 bytes of a `u64` slot.
    #[inline(always)]
    pub fn push_f32(&mut self, val: f32) {
        debug_assert!(self.len < MAX_ARGS, "RawArgs overflow");
        self.storage[self.len] = val.to_bits() as u64;
        self.len += 1;
    }

    /// Build the `void**` argument pointer array on the caller's stack.
    ///
    /// Each entry points to the corresponding `storage[i]` slot.
    /// The returned array is valid as long as `self` is not moved or
    /// dropped -- call this immediately before [`raw_launch`] and pass
    /// the result directly.
    #[inline(always)]
    pub fn build_arg_ptrs(&mut self) -> [*mut c_void; MAX_ARGS] {
        let mut ptrs = [std::ptr::null_mut(); MAX_ARGS];
        for i in 0..self.len {
            ptrs[i] = (&mut self.storage[i]) as *mut u64 as *mut c_void;
        }
        ptrs
    }
}

impl Default for RawArgs {
    fn default() -> Self {
        Self::new()
    }
}

/// Launch a CUDA kernel with pre-packed argument pointers.
///
/// Zero trait dispatch, zero Arc increment, zero event tracking.
/// The caller is responsible for stream synchronization and argument
/// validity.
///
/// # Safety
///
/// - `func` must be a valid, loaded `CUfunction` handle.
/// - `stream` must be a valid `CUstream` on the same device/context.
/// - `args` must contain exactly as many valid pointers as the kernel
///   expects, each pointing to a value of the correct type and size.
/// - The caller must ensure all device memory referenced by the kernel
///   remains valid until the kernel completes (async execution).
#[inline(always)]
pub unsafe fn raw_launch(
    func: CUfunction,
    grid: (u32, u32, u32),
    block: (u32, u32, u32),
    shared_mem_bytes: u32,
    stream: CUstream,
    args: &mut [*mut c_void],
) -> Result<(), CUresult> {
    let r = sys::cuLaunchKernel(
        func,
        grid.0, grid.1, grid.2,
        block.0, block.1, block.2,
        shared_mem_bytes,
        stream,
        args.as_mut_ptr(),
        std::ptr::null_mut(),
    );
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Launch a CUDA graph execution without cudarc's bind_to_thread or
/// event tracking overhead. In steady state the context is already
/// bound and synchronization is managed by stream ordering / events.
///
/// # Safety
///
/// - `graph_exec` must be a valid, instantiated `CUgraphExec`.
/// - `stream` must be a valid `CUstream` on the same device/context.
/// - The caller must ensure the graph's captured work is compatible
///   with the current device memory state.
#[inline(always)]
pub unsafe fn raw_graph_launch(
    graph_exec: CUgraphExec,
    stream: CUstream,
) -> Result<(), CUresult> {
    let r = sys::cuGraphLaunch(graph_exec, stream);
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Synchronize a CUDA stream without cudarc's bind_to_thread overhead.
///
/// # Safety
///
/// - `stream` must be a valid `CUstream`.
#[inline(always)]
pub unsafe fn raw_stream_sync(stream: CUstream) -> Result<(), CUresult> {
    let r = sys::cuStreamSynchronize(stream);
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Record a CUDA event on a stream without cudarc's bind_to_thread
/// or context equality checks.
///
/// # Safety
///
/// - `event` must be a valid `CUevent` on the same context as `stream`.
/// - `stream` must be a valid `CUstream`.
#[inline(always)]
pub unsafe fn raw_event_record(
    event: CUevent,
    stream: CUstream,
) -> Result<(), CUresult> {
    let r = sys::cuEventRecord(event, stream);
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Synchronize (host-wait) on a CUDA event without cudarc's
/// bind_to_thread overhead.
///
/// # Safety
///
/// - `event` must be a valid `CUevent`.
#[inline(always)]
pub unsafe fn raw_event_sync(event: CUevent) -> Result<(), CUresult> {
    let r = sys::cuEventSynchronize(event);
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Non-blocking query: has every command recorded before this event
/// completed? Returns `true` if done, `false` if still in-flight.
/// Any driver error other than `CUDA_ERROR_NOT_READY` is propagated.
///
/// # Safety
///
/// - `event` must be a valid `CUevent` that has been recorded via
///   `cuEventRecord`.
#[inline(always)]
pub unsafe fn raw_event_query(event: CUevent) -> Result<bool, CUresult> {
    let r = sys::cuEventQuery(event);
    match r {
        CUresult::CUDA_SUCCESS => Ok(true),
        CUresult::CUDA_ERROR_NOT_READY => Ok(false),
        other => Err(other),
    }
}

/// Make a stream wait on an event (GPU-side dependency) without
/// cudarc's bind_to_thread or context equality checks.
///
/// # Safety
///
/// - `stream` must be a valid `CUstream`.
/// - `event` must be a valid `CUevent` on the same context.
#[inline(always)]
pub unsafe fn raw_stream_wait_event(
    stream: CUstream,
    event: CUevent,
) -> Result<(), CUresult> {
    let r = sys::cuStreamWaitEvent(
        stream,
        event,
        sys::CUevent_wait_flags::CU_EVENT_WAIT_DEFAULT as u32,
    );
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Async memset (zero-fill) on device memory without cudarc's
/// bind_to_thread or event tracking overhead. Equivalent to
/// `cuMemsetD8Async(dst, 0, count, stream)`.
///
/// # Safety
///
/// - `dst` must be a valid device pointer with at least `count` bytes.
/// - `stream` must be a valid `CUstream` on the same device/context.
#[inline(always)]
pub unsafe fn raw_memset_d8_zero(
    dst: u64,
    count: usize,
    stream: CUstream,
) -> Result<(), CUresult> {
    let r = sys::cuMemsetD8Async(dst, 0, count, stream);
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Async device-to-device memcpy without cudarc's bind_to_thread or
/// event tracking overhead.
///
/// # Safety
///
/// - `dst` and `src` must be valid device pointers with at least
///   `count` bytes accessible.
/// - `stream` must be a valid `CUstream` on the same device/context.
#[inline(always)]
pub unsafe fn raw_memcpy_dtod_async(
    dst: u64,
    src: u64,
    count: usize,
    stream: CUstream,
) -> Result<(), CUresult> {
    let r = sys::cuMemcpyDtoDAsync_v2(dst, src, count, stream);
    if r != CUresult::CUDA_SUCCESS {
        return Err(r);
    }
    Ok(())
}

/// Extract the raw `CUfunction` handle from a cudarc `CudaFunction`.
#[inline(always)]
pub fn extract_fn(func: &cudarc::driver::CudaFunction) -> CUfunction {
    func.cu_function()
}

/// Extract the raw `CUstream` handle from a cudarc `CudaStream`.
#[inline(always)]
pub fn extract_stream(stream: &cudarc::driver::CudaStream) -> CUstream {
    stream.cu_stream()
}

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

    #[test]
    fn raw_args_push_and_build() {
        let mut args = RawArgs::new();
        assert!(args.is_empty());

        args.push_ptr(0xDEAD_BEEF_u64);
        args.push_i32(42);
        args.push_f32(3.14);
        args.push_u64(0xCAFE);

        assert_eq!(args.len(), 4);

        let ptrs = args.build_arg_ptrs();
        // Verify the first pointer points to storage[0] which holds 0xDEAD_BEEF
        unsafe {
            let val = *(ptrs[0] as *const u64);
            assert_eq!(val, 0xDEAD_BEEF_u64);
        }
        // Verify i32 stored correctly in low bytes
        unsafe {
            let val = *(ptrs[1] as *const u32);
            assert_eq!(val, 42);
        }
        // Verify f32 bit pattern
        unsafe {
            let val = *(ptrs[2] as *const u32);
            assert_eq!(f32::from_bits(val), 3.14);
        }
        // Verify u64
        unsafe {
            let val = *(ptrs[3] as *const u64);
            assert_eq!(val, 0xCAFE);
        }
        // Remaining pointers are null
        assert!(ptrs[4].is_null());
    }

    #[test]
    fn raw_args_reset_reuse() {
        let mut args = RawArgs::new();
        args.push_ptr(1);
        args.push_ptr(2);
        assert_eq!(args.len(), 2);

        args.reset();
        assert!(args.is_empty());
        assert_eq!(args.len(), 0);

        args.push_ptr(99);
        assert_eq!(args.len(), 1);
        let ptrs = args.build_arg_ptrs();
        unsafe {
            assert_eq!(*(ptrs[0] as *const u64), 99);
        }
    }

    #[test]
    fn push_i32_negative() {
        let mut args = RawArgs::new();
        args.push_i32(-1);
        let ptrs = args.build_arg_ptrs();
        // -1i32 as u32 = 0xFFFF_FFFF, CUDA reads 4 bytes and interprets as int
        unsafe {
            let val = *(ptrs[0] as *const i32);
            assert_eq!(val, -1);
        }
    }
}
