//! Mapped-pinned record buffers (CPU+GPU visible via `cuMemHostAlloc(DEVICEMAP)`).
//!
//! This is the **only** permitted CPU↔GPU path per
//! `feedback_no_htod_htoh_only_mapped_pinned.md`. Mirror of
//! `crates/ml/src/cuda_pipeline/mapped_pinned.rs::MappedF32Buffer` —
//! duplicated here because adding `ml-alpha -> ml` would cycle (ml
//! depends on ml-alpha for `mamba2_block`). Move both to
//! `ml-core::mapped_pinned` to deduplicate (deferred).

#![allow(unsafe_code)]

use std::ffi::c_void;
use std::mem::MaybeUninit;

/// Mapped-pinned host buffer of `T` records. Allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`; both host and device see the
/// same physical memory. Device pointer obtained via
/// `cuMemHostGetDevicePointer_v2`. The kernel writes via `dev_ptr`
/// (with `__threadfence_system()` in the kernel); the CPU reads via
/// `host_ptr` after stream synchronisation.
///
/// Per `feedback_no_htod_htoh_only_mapped_pinned`: the ONLY legitimate
/// channel for cross-direction transfer in the captured-graph hot path.
///
/// SAFETY: callers must ensure `T` is plain-old-data (no Drop, no
/// references, `repr(C)` recommended). The `Copy` bound enforces a
/// subset of this at compile time.
#[allow(missing_debug_implementations)]
pub struct MappedRecordBuffer<T: Copy> {
    pub host_ptr: *mut T,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    /// Number of `T`-records (NOT bytes).
    pub len: usize,
}

// Safety: only accessed from the thread that constructed it (which holds
// the active CUDA context). The host pointer is valid for the life of the
// allocation.
unsafe impl<T: Copy + Send> Send for MappedRecordBuffer<T> {}
unsafe impl<T: Copy + Sync> Sync for MappedRecordBuffer<T> {}

impl<T: Copy> MappedRecordBuffer<T> {
    /// Allocate `len` records of `T` as mapped-pinned memory.
    ///
    /// # Safety
    /// Caller must ensure a CUDA context is active on the current thread.
    pub unsafe fn new(len: usize) -> Result<Self, String> {
        let num_bytes = len * std::mem::size_of::<T>();
        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| format!("MappedF32Buffer alloc ({len} f32): {e}"))? as *mut T;

        // Zero-init so a stale read returns a safe default (all-zero bit pattern).
        std::ptr::write_bytes(host_ptr, 0, len);

        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| format!("cuMemHostGetDevicePointer (f32 buf): {e}"))?;

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

    /// Write a single record at index `idx`. Volatile write — the
    /// compiler cannot reorder the caller's store past this point.
    pub fn write_record(&self, idx: usize, val: T) {
        assert!(
            idx < self.len,
            "MappedRecordBuffer: idx {} >= len {}",
            idx,
            self.len
        );
        unsafe {
            std::ptr::write_volatile(self.host_ptr.add(idx), val);
        }
    }

    /// Read a single record at index `idx`. Volatile read — caller must
    /// have synchronised the producing stream first.
    pub fn read_record(&self, idx: usize) -> T {
        assert!(
            idx < self.len,
            "MappedRecordBuffer: idx {} >= len {}",
            idx,
            self.len
        );
        unsafe { std::ptr::read_volatile(self.host_ptr.add(idx)) }
    }

    /// Read all `len` entries via `read_volatile`. Caller must have
    /// synchronised the producing stream first.
    pub fn read_all(&self) -> Vec<T> {
        let mut out = Vec::with_capacity(self.len);
        unsafe {
            for i in 0..self.len {
                out.push(std::ptr::read_volatile(self.host_ptr.add(i)));
            }
        }
        out
    }

    /// Mutable host slice over the mapped pages.
    ///
    /// # Safety
    /// `host_ptr` is valid for `self.len` `T`s by construction; the
    /// caller must not invalidate the allocation while the slice is
    /// live and must respect unique-mutable-borrow semantics.
    pub fn host_slice_mut(&mut self) -> &mut [T] {
        unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) }
    }
}

impl<T: Copy> Drop for MappedRecordBuffer<T> {
    fn drop(&mut self) {
        unsafe {
            let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void);
        }
    }
}

/// Backward-compatible alias for the original `f32`-only buffer name.
/// All pre-existing call sites continue to work unchanged.
pub type MappedF32Buffer = MappedRecordBuffer<f32>;

impl MappedRecordBuffer<f32> {
    /// Write CPU-side data into the buffer via `host_ptr`. The kernel
    /// reading via `dev_ptr` sees the data after a stream sync barrier
    /// (mapped-pinned coherence).
    ///
    /// `f32`-specific bulk-write convenience for the legacy
    /// `MappedF32Buffer` call sites that pass `&[f32]`.
    pub fn write_from_slice(&self, slice: &[f32]) {
        assert!(
            slice.len() <= self.len,
            "MappedF32Buffer write overflow: have {} have {}",
            slice.len(),
            self.len
        );
        unsafe {
            for (i, &v) in slice.iter().enumerate() {
                std::ptr::write_volatile(self.host_ptr.add(i), v);
            }
        }
    }
}


/// Mapped-pinned `i32` buffer. Same allocation pattern as
/// [`MappedF32Buffer`] — the only permitted CPU↔GPU path per
/// `feedback_no_htod_htoh_only_mapped_pinned.md`. Needed by the
/// batched snap_feature path to stage per-snapshot `trade_count`
/// values for kernel consumption.
#[allow(missing_debug_implementations)]
pub struct MappedI32Buffer {
    pub host_ptr: *mut i32,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    pub len: usize,
}

unsafe impl Send for MappedI32Buffer {}
unsafe impl Sync for MappedI32Buffer {}

impl MappedI32Buffer {
    pub unsafe fn new(len: usize) -> Result<Self, String> {
        let num_bytes = len * std::mem::size_of::<i32>();
        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| format!("MappedI32Buffer alloc ({len} i32): {e}"))? as *mut i32;
        std::ptr::write_bytes(host_ptr, 0, len);
        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| format!("cuMemHostGetDevicePointer (i32 buf): {e}"))?;
        Ok(Self { host_ptr, dev_ptr: dev_ptr_raw.assume_init(), len })
    }
    pub fn host_slice_mut(&mut self) -> &mut [i32] {
        unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) }
    }
}
impl Drop for MappedI32Buffer {
    fn drop(&mut self) {
        unsafe { let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); }
    }
}


/// Mapped-pinned `i64` buffer. Needed by the batched snap_feature
/// path to stage per-snapshot `ts_ns` / `prev_ts_ns` timestamps.
#[allow(missing_debug_implementations)]
pub struct MappedI64Buffer {
    pub host_ptr: *mut i64,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    pub len: usize,
}

unsafe impl Send for MappedI64Buffer {}
unsafe impl Sync for MappedI64Buffer {}

impl MappedI64Buffer {
    pub unsafe fn new(len: usize) -> Result<Self, String> {
        let num_bytes = len * std::mem::size_of::<i64>();
        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| format!("MappedI64Buffer alloc ({len} i64): {e}"))? as *mut i64;
        std::ptr::write_bytes(host_ptr, 0, len);
        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| format!("cuMemHostGetDevicePointer (i64 buf): {e}"))?;
        Ok(Self { host_ptr, dev_ptr: dev_ptr_raw.assume_init(), len })
    }
    pub fn host_slice_mut(&mut self) -> &mut [i64] {
        unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) }
    }
}
impl Drop for MappedI64Buffer {
    fn drop(&mut self) {
        unsafe { let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void); }
    }
}
