//! Double-buffered GPU data loader for zero-downtime fold transitions.
//!
//! Holds two `DqnGpuData` slots (active + staging) so the next fold's data
//! can be uploaded to GPU while the current fold is still training.
//! Call `upload_to_staging()` then `swap()` to seamlessly transition.
//!
//! ## CUDA event-based async synchronization
//!
//! `upload_to_staging()` issues PCIe transfers on the CUDA stream, then records
//! a lightweight `CudaEvent` marking the upload's completion point.
//! `sync_staging()` polls the event via `is_complete()` (non-blocking) and only
//! falls back to `event.synchronize()` if the upload has not yet landed. This
//! eliminates the 2-4 ms host stall that a full `stream.synchronize()` caused
//! at every fold boundary.

use std::sync::Arc;
use cudarc::driver::{CudaEvent, CudaStream};
use tracing::info;

use super::DqnGpuData;
use crate::MLError;

/// Double-buffered GPU data loader.
///
/// While the trainer reads from `active()`, the next fold's data can be
/// uploaded into `staging` via `upload_to_staging()`.  Once the current
/// fold finishes, `swap()` promotes staging -> active in O(1).
///
/// `swap()` internally calls [`sync_staging()`](Self::sync_staging) to
/// ensure all GPU transfers are complete before the trainer reads the new data.
///
/// Synchronization uses a lightweight `CudaEvent` recorded after the upload
/// rather than a full `stream.synchronize()`.  The event is polled
/// non-blockingly first; only if the upload is still in flight does the host
/// block on `event.synchronize()`.
pub struct DoubleBufferedLoader {
    active: Option<DqnGpuData>,
    staging: Option<DqnGpuData>,
    stream: Arc<CudaStream>,
    /// Whether the staging upload has been synchronized (GPU transfers complete).
    /// Set to `false` by `upload_to_staging()`, set to `true` by `sync_staging()`.
    staging_synced: bool,
    /// CUDA event recorded on the stream immediately after the staging upload.
    /// `sync_staging()` waits on this event instead of synchronizing the entire
    /// stream.  `None` until the first staging upload.
    staging_event: Option<CudaEvent>,
}

impl std::fmt::Debug for DoubleBufferedLoader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut s = f.debug_struct("DoubleBufferedLoader");
        s.field("active", &self.active)
            .field("staging", &self.staging)
            .field("staging_synced", &self.staging_synced);
        s.field("staging_event", &self.staging_event.as_ref().map(|_| "CudaEvent(...)"));
        s.finish()
    }
}

impl DoubleBufferedLoader {
    /// Create a new double-buffered loader for the given CUDA stream.
    pub fn new(stream: Arc<CudaStream>) -> Self {
        Self {
            active: None,
            staging: None,
            stream,
            staging_synced: true, // no pending upload -> trivially synced
            staging_event: None,
        }
    }

    /// Upload initial data into the active slot.
    pub fn upload_initial(
        &mut self,
        data: &[([f64; 42], Vec<f64>)],
    ) -> Result<(), MLError> {
        let gpu_data = DqnGpuData::upload(data, &self.stream)?;
        info!(
            "DoubleBuffer: initial upload -- {} bars, {:.1} MB VRAM",
            gpu_data.num_bars,
            gpu_data.vram_bytes() as f64 / 1_048_576.0,
        );
        self.active = Some(gpu_data);
        Ok(())
    }

    /// Upload next fold's data into the staging slot (can overlap with training).
    ///
    /// The PCIe transfer is issued on the CUDA stream and a lightweight
    /// `CudaEvent` is recorded immediately after.  The data is *not* guaranteed
    /// to be resident until [`sync_staging()`](Self::sync_staging) (or
    /// [`swap()`](Self::swap), which calls it internally) completes.
    pub fn upload_to_staging(
        &mut self,
        data: &[([f64; 42], Vec<f64>)],
    ) -> Result<(), MLError> {
        let gpu_data = DqnGpuData::upload(data, &self.stream)?;
        info!(
            "DoubleBuffer: staging upload -- {} bars, {:.1} MB VRAM",
            gpu_data.num_bars,
            gpu_data.vram_bytes() as f64 / 1_048_576.0,
        );
        self.staging = Some(gpu_data);
        self.staging_synced = false; // upload in-flight until sync

        // Record a CudaEvent so sync_staging() can wait on just this upload
        // rather than synchronizing the entire stream.
        let event = self.stream.record_event(None).map_err(|e| {
            MLError::DeviceError(format!(
                "DoubleBuffer: staging event record failed: {e}"
            ))
        })?;
        self.staging_event = Some(event);

        Ok(())
    }

    /// Ensure the staging upload has landed on the device.
    ///
    /// First polls the staging event with `is_complete()` (non-blocking). If the
    /// upload has already finished (the common case -- training takes far longer
    /// than upload), this returns immediately without any host stall. Otherwise
    /// it falls back to `event.synchronize()`, which is still cheaper than a full
    /// `stream.synchronize()` because it only waits for work up to the recorded
    /// event.
    ///
    /// It is safe (but wasteful) to call this multiple times -- subsequent calls
    /// after the first are no-ops.
    ///
    /// [`swap()`](Self::swap) calls this internally, so callers only need to
    /// invoke it explicitly if they want to measure sync latency separately.
    pub fn sync_staging(&mut self) -> Result<(), MLError> {
        if self.staging_synced || self.staging.is_none() {
            return Ok(());
        }

        if let Some(ref event) = self.staging_event {
            // Fast path: non-blocking poll.  If the upload finished while
            // training was running (the typical case), we avoid any host
            // block at all.
            if !event.is_complete() {
                // Slow path: upload still in flight -- block until just
                // this event (not the whole stream) completes.
                event.synchronize().map_err(|e| {
                    MLError::DeviceError(format!(
                        "DoubleBuffer: staging event sync failed: {e}"
                    ))
                })?;
            }
        }

        self.staging_synced = true;
        Ok(())
    }

    /// Swap staging -> active.  Old active data is dropped (GPU memory freed).
    ///
    /// Calls [`sync_staging()`](Self::sync_staging) first to ensure the upload
    /// is fully resident on the device before the trainer reads from it.
    ///
    /// Returns `Err` if staging is empty (nothing was uploaded).
    pub fn swap(&mut self) -> Result<(), MLError> {
        // Ensure upload is complete before we hand the data to the trainer.
        self.sync_staging()?;

        let staging = self.staging.take().ok_or_else(|| {
            MLError::ModelError("DoubleBuffer: swap() called with empty staging slot".into())
        })?;
        let old_bars = self.active.as_ref().map_or(0, |d| d.num_bars);
        self.active = Some(staging);
        self.staging_synced = true; // staging is gone -> trivially synced
        info!(
            "DoubleBuffer: swapped (old {} bars -> new {} bars)",
            old_bars,
            self.active.as_ref().map_or(0, |d| d.num_bars),
        );
        Ok(())
    }

    /// Reference to the active (currently training) data.
    pub fn active(&self) -> Option<&DqnGpuData> {
        self.active.as_ref()
    }

    /// Whether staging data is ready to swap.
    pub fn staging_ready(&self) -> bool {
        self.staging.is_some()
    }

    /// Whether the staging upload has been synchronized on the device.
    ///
    /// Returns `true` if there is no pending upload or if `sync_staging()` has
    /// been called since the last `upload_to_staging()`.
    pub fn is_staging_synced(&self) -> bool {
        self.staging_synced
    }

    /// Total VRAM bytes used by both slots.
    pub fn total_vram_bytes(&self) -> usize {
        let a = self.active.as_ref().map_or(0, |d| d.vram_bytes());
        let s = self.staging.as_ref().map_or(0, |d| d.vram_bytes());
        a + s
    }

    /// Reference to the CUDA stream this loader uses.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }
}

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

    fn cuda_stream() -> Arc<CudaStream> {
        let context = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
        context.default_stream()
    }

    fn make_data(n: usize) -> Vec<([f64; 42], Vec<f64>)> {
        (0..n)
            .map(|i| {
                let features = [i as f64 * 0.01; 42];
                let targets = vec![100.0 + i as f64, 101.0, 100.5, 101.5];
                (features, targets)
            })
            .collect()
    }

    #[test]
    fn test_initial_upload() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        assert!(loader.active().is_none());

        loader.upload_initial(&make_data(50)).ok();
        // On CI without CUDA this will fail; on GPU it works
    }

    #[test]
    fn test_staging_and_swap() {
        let stream = cuda_stream();
        let mut loader = DoubleBufferedLoader::new(stream);
        if loader.upload_initial(&make_data(100)).is_err() {
            return; // skip on non-CUDA
        }
        if loader.upload_to_staging(&make_data(200)).is_err() {
            return;
        }

        assert!(loader.staging_ready());
        assert_eq!(loader.active().map_or(0, |d| d.num_bars), 100);

        loader.swap().ok();
        assert_eq!(loader.active().map_or(0, |d| d.num_bars), 200);
        assert!(!loader.staging_ready());
    }

    #[test]
    fn test_swap_without_staging_fails() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        if loader.upload_initial(&make_data(10)).is_err() {
            return;
        }
        assert!(loader.swap().is_err());
    }

    #[test]
    fn test_vram_tracking() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        if loader.upload_initial(&make_data(100)).is_err() {
            return;
        }
        let single = loader.total_vram_bytes();
        assert!(single > 0);

        if loader.upload_to_staging(&make_data(100)).is_err() {
            return;
        }
        let double = loader.total_vram_bytes();
        assert_eq!(double, 2 * single);
    }

    #[test]
    fn test_double_buffer_sync_staging() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        if loader.upload_initial(&make_data(50)).is_err() {
            return;
        }
        if loader.upload_to_staging(&make_data(100)).is_err() {
            return;
        }

        // After upload, staging is not yet synced
        assert!(!loader.is_staging_synced());

        // sync_staging uses event-based sync
        loader.sync_staging().ok();
        assert!(loader.is_staging_synced());
    }

    #[test]
    fn test_double_buffer_sync_staging_idempotent() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        if loader.upload_initial(&make_data(50)).is_err() {
            return;
        }
        if loader.upload_to_staging(&make_data(100)).is_err() {
            return;
        }

        // Multiple sync calls should all succeed
        loader.sync_staging().ok();
        loader.sync_staging().ok();
        assert!(loader.is_staging_synced());
    }

    #[test]
    fn test_double_buffer_sync_no_staging() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        // sync_staging with no staging data is a no-op
        loader.sync_staging().ok();
        assert!(loader.is_staging_synced());
    }

    #[test]
    fn test_double_buffer_swap_syncs_automatically() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        if loader.upload_initial(&make_data(50)).is_err() {
            return;
        }
        if loader.upload_to_staging(&make_data(100)).is_err() {
            return;
        }

        assert!(!loader.is_staging_synced());

        // swap() should call sync_staging internally
        loader.swap().ok();
        assert!(loader.is_staging_synced());
        assert_eq!(loader.active().map_or(0, |d| d.num_bars), 100);
    }

    #[test]
    fn test_double_buffer_synced_after_new() {
        let loader = DoubleBufferedLoader::new(cuda_stream());
        // Fresh loader has no pending uploads -> trivially synced
        assert!(loader.is_staging_synced());
    }

    #[test]
    fn test_double_buffer_upload_resets_synced() {
        let mut loader = DoubleBufferedLoader::new(cuda_stream());
        if loader.upload_initial(&make_data(50)).is_err() {
            return;
        }

        // First staging upload
        if loader.upload_to_staging(&make_data(100)).is_err() {
            return;
        }
        assert!(!loader.is_staging_synced());

        loader.sync_staging().ok();
        assert!(loader.is_staging_synced());

        // Second staging upload resets the flag
        if loader.upload_to_staging(&make_data(200)).is_err() {
            return;
        }
        assert!(!loader.is_staging_synced());
    }
}
