#![allow(unsafe_code)] // cuMemHostAlloc + cuMemHostGetDevicePointer_v2

//! Mapped pinned memory buffers for CPU↔GPU communication.
//!
//! `cuMemHostAlloc(DEVICEMAP | PORTABLE)` allocates host memory visible to the
//! GPU via a device pointer obtained through `cuMemHostGetDevicePointer_v2`.
//! Kernels write through `dev_ptr` (with `__threadfence_system()`); the host
//! reads back via `read_volatile` on `host_ptr` after stream synchronisation.
//!
//! This is the **only** permitted CPU↔GPU communication path per
//! `feedback_no_htod_htoh_only_mapped_pinned.md`. No `memcpy_dtoh` anywhere.
//!
//! Used by:
//! - `trainers/dqn/distributional_q_tests.rs` (Test 0.F kernel wrappers)
//! - Upcoming MoE Phase 2 kernel tests

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

// ── MappedI32Buffer ───────────────────────────────────────────────────────────

/// CPU+GPU visible buffer of `i32`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel writes via `dev_ptr`
/// (with `__threadfence_system()`); the CPU reads via `host_ptr`.
#[allow(missing_debug_implementations)] // raw pointer + CUdeviceptr have no useful Debug
pub struct MappedI32Buffer {
    pub host_ptr: *mut i32,
    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 MappedI32Buffer {}
unsafe impl Sync for MappedI32Buffer {}

impl MappedI32Buffer {
    /// Allocate `len` i32s 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::<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;

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

    /// Read all `len` entries via `read_volatile` (defeats CPU caching of
    /// stale values written before the kernel completed). Caller must have
    /// synchronised the producing stream first.
    pub fn read_all(&self) -> Vec<i32> {
        let mut out = Vec::with_capacity(self.len);
        // Safety: host_ptr is valid for self.len i32s by construction.
        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.
    /// Direct memory write — no memcpy. The kernel reading via dev_ptr
    /// sees the data after a stream sync barrier (mapped pinned coherence).
    /// Caller must ensure `slice.len() <= self.len`.
    pub fn write_from_slice(&self, slice: &[i32]) {
        assert!(slice.len() <= self.len, "MappedI32Buffer write overflow");
        // Safety: host_ptr is valid for self.len i32s by construction.
        unsafe {
            for (i, &v) in slice.iter().enumerate() {
                std::ptr::write_volatile(self.host_ptr.add(i), v);
            }
        }
    }
}

impl Drop for MappedI32Buffer {
    fn drop(&mut self) {
        // Safety: host_ptr was returned by malloc_host and is non-null.
        // Best-effort free — nothing actionable on failure inside Drop.
        unsafe {
            #[allow(clippy::let_underscore_must_use)]
            let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void);
        }
    }
}

// ── MappedU32Buffer ───────────────────────────────────────────────────────────

/// CPU+GPU visible buffer of `u32`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. Kernels read/write via `dev_ptr`;
/// the CPU reads via `read_volatile` on `host_ptr`.
pub struct MappedU32Buffer {
    pub host_ptr: *mut u32,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    pub len: usize,
}

impl std::fmt::Debug for MappedU32Buffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MappedU32Buffer")
            .field("dev_ptr", &self.dev_ptr)
            .field("len", &self.len)
            .finish()
    }
}

unsafe impl Send for MappedU32Buffer {}
unsafe impl Sync for MappedU32Buffer {}

impl MappedU32Buffer {
    /// Allocate `len` u32s 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::<u32>();
        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!("MappedU32Buffer alloc ({len} u32): {e}"))?
            as *mut u32;

        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 (u32 buf): {e}"))?;

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

    /// Write CPU-side data into the buffer via host_ptr (no memcpy).
    /// Caller must ensure `slice.len() <= self.len`.
    pub fn write_from_slice(&self, slice: &[u32]) {
        assert!(slice.len() <= self.len, "MappedU32Buffer write overflow");
        // Safety: host_ptr is valid for self.len u32s by construction.
        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 — `&mut [u32]` aliasing
    /// `host_ptr..host_ptr+len`. Lets callers structure writes without
    /// per-element `write_volatile`. The kernel sees the writes after a
    /// stream-sync barrier.
    ///
    /// # Safety
    /// `host_ptr` is valid for `self.len` u32s by construction; the caller
    /// must respect the unique-mutable-borrow rule for `&mut [T]`.
    pub fn host_slice_mut(&mut self) -> &mut [u32] {
        // Safety: see invariants above.
        unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) }
    }
}

impl Drop for MappedU32Buffer {
    fn drop(&mut self) {
        unsafe {
            #[allow(clippy::let_underscore_must_use)]
            let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void);
        }
    }
}

// ── MappedF32Buffer ───────────────────────────────────────────────────────────

/// CPU+GPU visible buffer of `f32`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel writes via `dev_ptr`
/// (with `__threadfence_system()`); the CPU reads via `host_ptr`.
#[allow(missing_debug_implementations)] // raw pointer + CUdeviceptr have no useful Debug
pub struct MappedF32Buffer {
    pub host_ptr: *mut f32,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    pub len: usize,
}

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);
        // Safety: host_ptr is valid for self.len f32s by construction.
        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.
    /// Direct memory write — no memcpy. The kernel reading via dev_ptr
    /// sees the data after a stream sync barrier (mapped pinned coherence).
    /// Caller must ensure `slice.len() <= self.len`.
    pub fn write_from_slice(&self, slice: &[f32]) {
        assert!(slice.len() <= self.len, "MappedF32Buffer write overflow");
        // Safety: host_ptr is valid for self.len f32s by construction.
        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 — `&mut [f32]` aliasing
    /// `host_ptr..host_ptr+len`. Lets callers structure writes (e.g.
    /// per-episode portfolio init) without per-element `write_volatile`.
    /// The kernel sees the writes after a stream-sync barrier (mapped
    /// pinned coherence — same memory, different alias).
    ///
    /// # 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 the unique-mutable-borrow rule for `&mut [T]`.
    pub fn host_slice_mut(&mut self) -> &mut [f32] {
        // Safety: see invariants above.
        unsafe { std::slice::from_raw_parts_mut(self.host_ptr, self.len) }
    }
}

impl Drop for MappedF32Buffer {
    fn drop(&mut self) {
        // Safety: host_ptr was returned by malloc_host and is non-null.
        unsafe {
            #[allow(clippy::let_underscore_must_use)]
            let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void);
        }
    }
}


// ── MappedU64Buffer ───────────────────────────────────────────────────────────

/// CPU+GPU visible buffer of `u64`s allocated via
/// `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. Used for descriptor tables (e.g.
/// spectral-norm pointer/dim arrays) where each entry is a 64-bit value.
pub struct MappedU64Buffer {
    pub host_ptr: *mut u64,
    pub dev_ptr: cudarc::driver::sys::CUdeviceptr,
    pub len: usize,
}

impl std::fmt::Debug for MappedU64Buffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MappedU64Buffer")
            .field("dev_ptr", &self.dev_ptr)
            .field("len", &self.len)
            .finish()
    }
}

unsafe impl Send for MappedU64Buffer {}
unsafe impl Sync for MappedU64Buffer {}

impl MappedU64Buffer {
    /// Allocate `len` u64s 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::<u64>();
        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!("MappedU64Buffer alloc ({len} u64): {e}"))?
            as *mut u64;

        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 (u64 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<u64> {
        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.
    /// Direct memory write — no memcpy. Caller must ensure
    /// `slice.len() <= self.len`.
    pub fn write_from_slice(&self, slice: &[u64]) {
        assert!(slice.len() <= self.len, "MappedU64Buffer write overflow");
        unsafe {
            for (i, &v) in slice.iter().enumerate() {
                std::ptr::write_volatile(self.host_ptr.add(i), v);
            }
        }
    }
}

impl Drop for MappedU64Buffer {
    fn drop(&mut self) {
        unsafe {
            #[allow(clippy::let_underscore_must_use)]
            let _ = cudarc::driver::result::free_host(self.host_ptr as *mut c_void);
        }
    }
}
