//! Async non-blocking diagnostic staging via separate CUDA stream +
//! double-buffered mapped-pinned memory.
//!
//! The training loop reads ~7 device buffers to host every step for JSONL
//! diagnostics. Each blocking `read_slice_d` does `stream.synchronize()`
//! which stalls the GPU training pipeline. At b=256 this wastes ~5ms/step.
//!
//! `DiagStaging` solves this by:
//! 1. Running DtoD copies on a **separate** CUDA stream (non-blocking on
//!    the training stream).
//! 2. Using **double-buffered** mapped-pinned memory: the host reads the
//!    PREVIOUS step's buffer while the current step's copy runs on the
//!    diag stream.
//!
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: mapped-pinned is the
//! only permitted CPU<->GPU path. The DtoD copies land in mapped-pinned
//! device addresses; after `diag_stream.synchronize()` the host reads
//! the same physical pages via host_ptr.
//!
//! One-step latency: step 0 reports zeros (no prior snapshot). From step 1
//! onward, the host reads step N-1's data. Acceptable for diagnostics.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::driver::{CudaEvent, CudaStream};
use cudarc::driver::sys::CUstream;
use ml_core::device::MlDevice;

use crate::pinned_mem::MappedF32Buffer;
use crate::rl::frd::FRD_OUT_DIM;
use crate::rl::isv_slots::RL_SLOTS_END;
use crate::trainer::raw_launch::{raw_memcpy_dtod_async, raw_event_record, raw_event_sync};

/// Double-buffered mapped-pinned staging for async diagnostic reads.
///
/// Layout within each buffer (all f32, contiguous):
///   `[ISV (RL_SLOTS_END) | rewards (B) | dones (B) | actions (B) | raw_rewards (B)
///    | trade_duration (B) | outcome_ema (B) | position_lots (B) | pyramid_count (B)
///    | unit_entry_price (B×4) | unit_entry_step (B×4) | unit_lots (B×4)
///    | unit_trail (B×4) | close_unit_index (B) | frd_logits (B×FRD_OUT_DIM)]`
///
/// i32 fields (`actions`, `position_lots`, `pyramid_count`, `unit_entry_step`,
/// `unit_lots`, `close_unit_index`) are stored as raw bit patterns in f32 slots.
/// DtoD copies raw bytes (4 bytes per element, same width); host reinterprets
/// via `f32::to_bits() as i32`.
pub struct DiagStaging {
    diag_stream: Arc<CudaStream>,
    /// Cached raw `CUstream` handle for the diag stream — avoids
    /// cudarc's `bind_to_thread` overhead on every memcpy/sync call.
    raw_diag_stream: CUstream,
    /// Event recorded after each `snapshot_async` call. `sync_and_swap`
    /// syncs on this event instead of host-blocking the entire diag
    /// stream (`cuStreamSynchronize`). By the time `sync_and_swap` runs
    /// (after the main training step), the DtoD copies have long
    /// completed — the event sync is instant (~0 us).
    snapshot_done_event: CudaEvent,
    bufs: [MappedF32Buffer; 2],
    current: usize,
    /// Total f32 count per buffer.
    pub n_floats: usize,
    // Layout offsets (all in f32 elements from start of buffer).
    pub isv_offset: usize,
    pub isv_len: usize,
    pub rewards_offset: usize,
    pub dones_offset: usize,
    pub actions_offset: usize,
    pub raw_rewards_offset: usize,
    pub trade_duration_offset: usize,
    pub outcome_ema_offset: usize,
    pub position_lots_offset: usize,
    pub pyramid_count_offset: usize,
    pub unit_entry_price_offset: usize,
    pub unit_entry_step_offset: usize,
    pub unit_lots_offset: usize,
    pub unit_trail_offset: usize,
    pub close_unit_index_offset: usize,
    pub frd_logits_offset: usize,
    pub b_size: usize,
}

impl DiagStaging {
    /// Allocate double-buffered mapped-pinned staging for `b_size` batch slots.
    ///
    /// Creates a new CUDA stream on the same context as `dev`. The stream is
    /// independent of the training stream — copies launched on it do not stall
    /// the training pipeline.
    pub fn new(dev: &MlDevice, b_size: usize) -> Result<Self> {
        let ctx = dev.cuda_context().context("DiagStaging needs CUDA context")?;
        let diag_stream = ctx
            .new_stream()
            .map_err(|e| anyhow::anyhow!("DiagStaging new_stream: {e}"))?;

        // Layout (f32 slots):
        //   ISV | rewards | dones | actions | raw_rewards | trade_duration | outcome_ema
        //   | position_lots | pyramid_count | unit_entry_price(B×4) | unit_entry_step(B×4)
        //   | unit_lots(B×4) | unit_trail(B×4) | close_unit_index | frd_logits(B×FRD_OUT_DIM)
        let isv_len = RL_SLOTS_END;
        let mut off = isv_len;
        let rewards_offset = off; off += b_size;
        let dones_offset = off; off += b_size;
        let actions_offset = off; off += b_size;
        let raw_rewards_offset = off; off += b_size;
        let trade_duration_offset = off; off += b_size;
        let outcome_ema_offset = off; off += b_size;
        let position_lots_offset = off; off += b_size;
        let pyramid_count_offset = off; off += b_size;
        let unit_entry_price_offset = off; off += b_size * 4;
        let unit_entry_step_offset = off; off += b_size * 4;
        let unit_lots_offset = off; off += b_size * 4;
        let unit_trail_offset = off; off += b_size * 4;
        let close_unit_index_offset = off; off += b_size;
        let frd_logits_offset = off; off += b_size * FRD_OUT_DIM;
        let n_floats = off;

        let buf0 = unsafe { MappedF32Buffer::new(n_floats) }
            .map_err(|e| anyhow::anyhow!("DiagStaging buf0: {e}"))?;
        let buf1 = unsafe { MappedF32Buffer::new(n_floats) }
            .map_err(|e| anyhow::anyhow!("DiagStaging buf1: {e}"))?;

        let raw_diag_stream = diag_stream.cu_stream();
        let snapshot_done_event = ctx
            .new_event(None)
            .map_err(|e| anyhow::anyhow!("DiagStaging snapshot_done_event: {e}"))?;
        Ok(Self {
            diag_stream,
            raw_diag_stream,
            snapshot_done_event,
            bufs: [buf0, buf1],
            current: 0,
            n_floats,
            isv_offset: 0,
            isv_len,
            rewards_offset,
            dones_offset,
            actions_offset,
            raw_rewards_offset,
            trade_duration_offset,
            outcome_ema_offset,
            position_lots_offset,
            pyramid_count_offset,
            unit_entry_price_offset,
            unit_entry_step_offset,
            unit_lots_offset,
            unit_trail_offset,
            close_unit_index_offset,
            frd_logits_offset,
            b_size,
        })
    }

    /// Launch async DtoD copies from device buffers into the CURRENT staging
    /// buffer. Non-blocking on the training stream -- runs on the diag_stream.
    ///
    /// Accepts pre-extracted raw `u64` device pointers (`CudaSlice::raw_ptr()`)
    /// instead of `&CudaSlice` references. This bypasses cudarc's
    /// `device_ptr()` event-tracking path which internally does
    /// `cuStreamWaitEvent` + `cuEventRecord` per call -- eliminating 15
    /// inter-stream sync events per step.
    ///
    /// Safety contract: callers must pass pointers from `CudaSlice` fields
    /// whose device addresses are stable for the trainer's lifetime (all
    /// pre-allocated at init). The training stream must have completed
    /// writing to these buffers before the diag_stream reads them; this is
    /// guaranteed by CUDA's same-context stream ordering when the training
    /// stream's prior kernel launches have been enqueued before this call.
    #[allow(clippy::too_many_arguments)]
    pub fn snapshot_async(
        &mut self,
        isv_ptr: u64,
        rewards_ptr: u64,
        dones_ptr: u64,
        actions_ptr: u64,
        raw_rewards_ptr: u64,
        trade_duration_ptr: u64,
        outcome_ema_ptr: u64,
        position_lots_ptr: u64,
        pyramid_count_ptr: u64,
        unit_entry_price_ptr: u64,
        unit_entry_step_ptr: u64,
        unit_lots_ptr: u64,
        unit_trail_ptr: u64,
        close_unit_index_ptr: u64,
        frd_logits_ptr: u64,
    ) -> Result<()> {
        let buf = &self.bufs[self.current];
        let s = self.raw_diag_stream;
        let b = self.b_size;

        unsafe {
            // ISV
            raw_memcpy_dtod_async(
                buf.dev_ptr,
                isv_ptr,
                self.isv_len * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot isv: {:?}", e))?;

            // rewards
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.rewards_offset * 4) as u64,
                rewards_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot rewards: {:?}", e))?;

            // dones
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.dones_offset * 4) as u64,
                dones_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot dones: {:?}", e))?;

            // actions (i32 -- same byte width as f32, reinterpret on host)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.actions_offset * 4) as u64,
                actions_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot actions: {:?}", e))?;

            // raw_rewards
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.raw_rewards_offset * 4) as u64,
                raw_rewards_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot raw_rewards: {:?}", e))?;

            // trade_duration
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.trade_duration_offset * 4) as u64,
                trade_duration_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot trade_duration: {:?}", e))?;

            // outcome_ema
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.outcome_ema_offset * 4) as u64,
                outcome_ema_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot outcome_ema: {:?}", e))?;

            // position_lots (i32 -- raw bytes into f32 slot)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.position_lots_offset * 4) as u64,
                position_lots_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot position_lots: {:?}", e))?;

            // pyramid_count (i32 -- raw bytes into f32 slot)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.pyramid_count_offset * 4) as u64,
                pyramid_count_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot pyramid_count: {:?}", e))?;

            // unit_entry_price (f32, B x 4)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.unit_entry_price_offset * 4) as u64,
                unit_entry_price_ptr,
                b * 4 * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot unit_entry_price: {:?}", e))?;

            // unit_entry_step (i32, B x 4 -- raw bytes)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.unit_entry_step_offset * 4) as u64,
                unit_entry_step_ptr,
                b * 4 * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot unit_entry_step: {:?}", e))?;

            // unit_lots (i32, B x 4 -- raw bytes)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.unit_lots_offset * 4) as u64,
                unit_lots_ptr,
                b * 4 * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot unit_lots: {:?}", e))?;

            // unit_trail (f32, B x 4)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.unit_trail_offset * 4) as u64,
                unit_trail_ptr,
                b * 4 * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot unit_trail: {:?}", e))?;

            // close_unit_index (i32, B -- raw bytes)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.close_unit_index_offset * 4) as u64,
                close_unit_index_ptr,
                b * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot close_unit_index: {:?}", e))?;

            // frd_logits (f32, B x FRD_OUT_DIM)
            raw_memcpy_dtod_async(
                buf.dev_ptr + (self.frd_logits_offset * 4) as u64,
                frd_logits_ptr,
                b * FRD_OUT_DIM * 4,
                s,
            )
            .map_err(|e| anyhow::anyhow!("diag snapshot frd_logits: {:?}", e))?;

            // Record event after all DtoD copies. sync_and_swap syncs
            // on this event instead of blocking the entire diag stream.
            raw_event_record(self.snapshot_done_event.cu_event(), s)
                .map_err(|e| anyhow::anyhow!("diag snapshot_done_event record: {:?}", e))?;
        }

        Ok(())
    }

    /// Wait for the previous `snapshot_async`'s DtoD copies to complete,
    /// then swap buffers. Uses event sync instead of full stream sync —
    /// by the time this runs (after the main training step), the DtoD
    /// copies have long completed, so the event sync is instant (~0 us).
    pub fn sync_and_swap(&mut self) -> Result<()> {
        unsafe {
            raw_event_sync(self.snapshot_done_event.cu_event())
                .map_err(|e| anyhow::anyhow!("diag snapshot_done_event sync: {:?}", e))?;
        }
        self.current = 1 - self.current;
        Ok(())
    }

    /// The owned `CudaStream` backing the raw handle. Kept alive for the
    /// lifetime of `DiagStaging`; dropping it would invalidate `raw_diag_stream`.
    pub fn diag_stream(&self) -> &CudaStream {
        &self.diag_stream
    }

    /// Read the PREVIOUS step's ISV data (the buffer NOT currently targeted
    /// for writes). Call AFTER `sync_and_swap`.
    pub fn read_isv(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.isv_offset), self.isv_len) }
    }

    /// Read the PREVIOUS step's per-batch rewards.
    pub fn read_rewards(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.rewards_offset), self.b_size) }
    }

    /// Read the PREVIOUS step's per-batch done flags.
    pub fn read_dones(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.dones_offset), self.b_size) }
    }

    /// Read the PREVIOUS step's per-batch actions as raw f32 bit patterns.
    /// The underlying data is i32; caller reinterprets via `f32::to_bits() as i32`.
    pub fn read_actions_raw(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.actions_offset), self.b_size) }
    }

    /// Read the PREVIOUS step's per-batch raw (pre-scale) rewards.
    pub fn read_raw_rewards(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe {
            std::slice::from_raw_parts(buf.host_ptr.add(self.raw_rewards_offset), self.b_size)
        }
    }

    /// Read the PREVIOUS step's per-batch trade duration.
    pub fn read_trade_duration(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe {
            std::slice::from_raw_parts(buf.host_ptr.add(self.trade_duration_offset), self.b_size)
        }
    }

    /// Read the PREVIOUS step's per-batch outcome EMA.
    pub fn read_outcome_ema(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe {
            std::slice::from_raw_parts(buf.host_ptr.add(self.outcome_ema_offset), self.b_size)
        }
    }

    /// Read the PREVIOUS step's per-batch position lots as i32.
    /// The underlying data is i32 stored as raw bit patterns in f32 slots.
    pub fn read_position_lots(&self) -> Vec<i32> {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        let raw = unsafe {
            std::slice::from_raw_parts(buf.host_ptr.add(self.position_lots_offset), self.b_size)
        };
        raw.iter().map(|f| f.to_bits() as i32).collect()
    }

    /// Read the PREVIOUS step's per-batch pyramid unit count as i32.
    pub fn read_pyramid_count(&self) -> Vec<i32> {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        let raw = unsafe {
            std::slice::from_raw_parts(buf.host_ptr.add(self.pyramid_count_offset), self.b_size)
        };
        raw.iter().map(|f| f.to_bits() as i32).collect()
    }

    /// Read the PREVIOUS step's per-batch per-unit entry prices (B x 4 f32).
    pub fn read_unit_entry_price(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe {
            std::slice::from_raw_parts(
                buf.host_ptr.add(self.unit_entry_price_offset),
                self.b_size * 4,
            )
        }
    }

    /// Read the PREVIOUS step's per-batch per-unit entry steps as i32 (B x 4).
    pub fn read_unit_entry_step(&self) -> Vec<i32> {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        let raw = unsafe {
            std::slice::from_raw_parts(
                buf.host_ptr.add(self.unit_entry_step_offset),
                self.b_size * 4,
            )
        };
        raw.iter().map(|f| f.to_bits() as i32).collect()
    }

    /// Read the PREVIOUS step's per-batch per-unit lot counts as i32 (B x 4).
    pub fn read_unit_lots(&self) -> Vec<i32> {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        let raw = unsafe {
            std::slice::from_raw_parts(
                buf.host_ptr.add(self.unit_lots_offset),
                self.b_size * 4,
            )
        };
        raw.iter().map(|f| f.to_bits() as i32).collect()
    }

    /// Read the PREVIOUS step's per-batch per-unit trail distances (B x 4 f32).
    pub fn read_unit_trail(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe {
            std::slice::from_raw_parts(
                buf.host_ptr.add(self.unit_trail_offset),
                self.b_size * 4,
            )
        }
    }

    /// Read the PREVIOUS step's per-batch close unit index as i32.
    pub fn read_close_unit_index(&self) -> Vec<i32> {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        let raw = unsafe {
            std::slice::from_raw_parts(
                buf.host_ptr.add(self.close_unit_index_offset),
                self.b_size,
            )
        };
        raw.iter().map(|f| f.to_bits() as i32).collect()
    }

    /// Read the PREVIOUS step's per-batch FRD logits (B x FRD_OUT_DIM f32).
    pub fn read_frd_logits(&self) -> &[f32] {
        let prev = 1 - self.current;
        let buf = &self.bufs[prev];
        unsafe {
            std::slice::from_raw_parts(
                buf.host_ptr.add(self.frd_logits_offset),
                self.b_size * FRD_OUT_DIM,
            )
        }
    }
}
