//! Mapped-pinned f32 buffer (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;

/// CPU+GPU visible buffer of `f32`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel writes via `dev_ptr`
/// (with `__threadfence_system()` in the kernel); the CPU reads via
/// `host_ptr` after stream synchronisation.
#[allow(missing_debug_implementations)]
pub struct MappedF32Buffer {
    pub host_ptr: *mut f32,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    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 Send for MappedF32Buffer {}
unsafe impl Sync for MappedF32Buffer {}

impl MappedF32Buffer {
    /// Allocate `len` f32s of 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::<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| format!("MappedF32Buffer alloc ({len} f32): {e}"))? as *mut f32;

        // Zero-init so a stale read returns a safe default (0.0).
        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,
        })
    }

    /// Read all `len` entries via `read_volatile`. Caller must have
    /// synchronised the producing stream first.
    pub fn read_all(&self) -> Vec<f32> {
        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
    }

    /// 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).
    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);
            }
        }
    }

    /// Mutable host slice over the mapped pages.
    ///
    /// # Safety
    /// `host_ptr` is valid for `self.len` f32s 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 [f32] {
        unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) }
    }
}

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


/// 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); }
    }
}
