//! Multi-GPU FFT with 1D slab decomposition.
//!
//! Distributes large FFTs across multiple GPUs using a 1D slab
//! decomposition strategy. Each GPU receives a contiguous slab of
//! the input data and performs local FFTs, followed by a global
//! transpose (all-to-all redistribution) and a second round of
//! local FFTs.
//!
//! ## Algorithm
//!
//! The multi-GPU FFT follows a three-phase approach:
//!
//! 1. **Local FFT (Phase 1):** Each device computes FFTs on its N/P elements.
//! 2. **Global transpose:** All-to-all redistribution via peer copy or host staging.
//! 3. **Local FFT (Phase 2):** Each device computes FFTs on the transposed slab.

use std::fmt;
use std::time::{Duration, Instant};

use oxicuda_driver::ffi::CUdeviceptr;
use oxicuda_memory::peer_copy;
use oxicuda_memory::{DeviceBuffer, PinnedBuffer, copy};

use crate::error::{FftError, FftResult};
use crate::execute::FftHandle;
use crate::plan::FftPlan;
use crate::types::{FftDirection, FftPrecision, FftType};

// ---------------------------------------------------------------------------
// TransposeRegion -- describes a data region to transfer between devices
// ---------------------------------------------------------------------------

/// Describes a contiguous region of elements to transfer between two GPUs
/// during the global transpose phase.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TransposeRegion {
    /// Offset (in elements) within the source device's slab.
    pub src_offset: usize,
    /// Offset (in elements) within the destination device's slab.
    pub dst_offset: usize,
    /// Number of elements to transfer.
    pub count: usize,
}

// ---------------------------------------------------------------------------
// SlabDecomposition -- partitions N elements across P devices
// ---------------------------------------------------------------------------

/// 1D slab decomposition that partitions `total_n` elements across
/// `num_devices` GPUs.
///
/// When `total_n` is not evenly divisible by `num_devices`, the first
/// `total_n % num_devices` devices receive one extra element each.
#[derive(Debug, Clone)]
pub struct SlabDecomposition {
    /// Total number of elements in the FFT.
    pub total_n: usize,
    /// Number of GPU devices participating.
    pub num_devices: usize,
    /// Per-device element counts (length == `num_devices`).
    pub slab_sizes: Vec<usize>,
    /// Cumulative offsets into the global array (length == `num_devices`).
    pub slab_offsets: Vec<usize>,
}

impl SlabDecomposition {
    /// Creates a new slab decomposition.
    ///
    /// # Errors
    ///
    /// Returns [`FftError::InvalidSize`] if `total_n` is zero or
    /// `num_devices` is zero, or if `total_n < num_devices`.
    pub fn new(total_n: usize, num_devices: usize) -> FftResult<Self> {
        if total_n == 0 {
            return Err(FftError::InvalidSize(
                "total_n must be > 0 for slab decomposition".to_string(),
            ));
        }
        if num_devices == 0 {
            return Err(FftError::InvalidSize(
                "num_devices must be > 0 for slab decomposition".to_string(),
            ));
        }
        if total_n < num_devices {
            return Err(FftError::InvalidSize(format!(
                "total_n ({total_n}) must be >= num_devices ({num_devices})"
            )));
        }

        let base_size = total_n / num_devices;
        let remainder = total_n % num_devices;

        let mut slab_sizes = Vec::with_capacity(num_devices);
        let mut slab_offsets = Vec::with_capacity(num_devices);
        let mut offset = 0usize;

        for i in 0..num_devices {
            slab_offsets.push(offset);
            let size = if i < remainder {
                base_size + 1
            } else {
                base_size
            };
            slab_sizes.push(size);
            offset += size;
        }

        Ok(Self {
            total_n,
            num_devices,
            slab_sizes,
            slab_offsets,
        })
    }

    /// Returns the number of elements assigned to `device_idx`.
    ///
    /// Returns 0 if `device_idx` is out of range.
    pub fn slab_size(&self, device_idx: usize) -> usize {
        self.slab_sizes.get(device_idx).copied().unwrap_or(0)
    }

    /// Returns the global offset for the slab assigned to `device_idx`.
    ///
    /// Returns 0 if `device_idx` is out of range.
    pub fn slab_offset(&self, device_idx: usize) -> usize {
        self.slab_offsets.get(device_idx).copied().unwrap_or(0)
    }

    /// Computes the transpose region describing which elements device
    /// `src_device` must send to device `dst_device` during the
    /// all-to-all transpose phase.
    ///
    /// In a 1D slab decomposition, the transpose redistributes data so
    /// that each device ends up with a different contiguous slice of the
    /// frequency domain. Device `src_device` sends its portion that
    /// overlaps with `dst_device`'s target range.
    ///
    /// Returns `None` if either device index is out of range.
    pub fn transpose_map(&self, src_device: usize, dst_device: usize) -> Option<TransposeRegion> {
        if src_device >= self.num_devices || dst_device >= self.num_devices {
            return None;
        }

        // In the transpose, each device's slab of size S_i gets split
        // into num_devices sub-blocks. The sub-block destined for
        // dst_device has size proportional to dst_device's slab size.
        //
        // For a uniform distribution: each src sends
        //   slab_size(dst) elements starting at offset slab_offset(dst)
        //   within the src's local buffer (relative to the global array).
        //
        // More precisely, src_device owns global indices
        //   [src_off, src_off + src_size)
        // and dst_device's target range in the transposed layout is
        //   [dst_off, dst_off + dst_size)
        //
        // The intersection determines what to send.
        let src_off = self.slab_offsets[src_device];
        let src_size = self.slab_sizes[src_device];
        let dst_off = self.slab_offsets[dst_device];
        let dst_size = self.slab_sizes[dst_device];

        // Compute intersection of [src_off, src_off+src_size) and
        // [dst_off, dst_off+dst_size)
        let inter_start = src_off.max(dst_off);
        let inter_end = (src_off + src_size).min(dst_off + dst_size);

        if inter_start >= inter_end {
            // No overlap -- for the general transpose, we still need to
            // describe the sub-block. In a full slab decomposition transpose,
            // each device sends a chunk of its local data to every other device.
            // The chunk for dst_device is: elements at positions corresponding
            // to dst_device's "column range" in the 2D interpretation.
            //
            // For a 1D FFT of size N distributed across P devices:
            // - Phase 1 computes local FFTs of size S_i on each device
            // - The transpose reinterprets the data as a 2D (P x S) array
            //   and transposes it
            //
            // Simplified model: src sends dst_size elements from its local
            // buffer, with src_offset = dst_off within the local slab,
            // and dst_offset = src_off within dst's buffer.
            let count = compute_transpose_count(src_size, dst_size, self.total_n);

            Some(TransposeRegion {
                src_offset: dst_device * (src_size / self.num_devices),
                dst_offset: src_device * (dst_size / self.num_devices),
                count,
            })
        } else {
            let count = inter_end - inter_start;
            Some(TransposeRegion {
                src_offset: inter_start - src_off,
                dst_offset: inter_start - dst_off,
                count,
            })
        }
    }
}

/// Computes the number of elements to transfer between two devices
/// in the transpose phase when their slabs do not overlap.
fn compute_transpose_count(src_size: usize, dst_size: usize, total_n: usize) -> usize {
    // Each device pair exchanges (src_size * dst_size) / total_n elements
    // (rounded up to avoid losing data).
    let numerator = src_size * dst_size;
    numerator.div_ceil(total_n)
}

impl fmt::Display for SlabDecomposition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "SlabDecomposition(N={}, P={}, sizes={:?})",
            self.total_n, self.num_devices, self.slab_sizes
        )
    }
}

// ---------------------------------------------------------------------------
// TransposeStrategy -- how to move data between devices
// ---------------------------------------------------------------------------

/// Strategy for executing the global transpose between GPUs.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransposeStrategy {
    /// Direct GPU-to-GPU copy using `copy_peer_async`.
    /// Best when NVLink or PCIe P2P access is available.
    PeerToPeer,
    /// GPU -> Host -> GPU copy, used when direct P2P is unavailable
    /// (e.g. across NUMA nodes without NVLink).
    StagedViaHost,
}

impl TransposeStrategy {
    /// Automatically selects the best transpose strategy based on
    /// the number of devices.
    ///
    /// For 2 devices, P2P is almost always available and efficient.
    /// For larger device counts, staged-via-host is safer as a default
    /// since not all device pairs may support P2P.
    pub fn auto_select(num_devices: usize) -> Self {
        if num_devices <= 2 {
            Self::PeerToPeer
        } else {
            Self::StagedViaHost
        }
    }
}

impl fmt::Display for TransposeStrategy {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::PeerToPeer => write!(f, "PeerToPeer"),
            Self::StagedViaHost => write!(f, "StagedViaHost"),
        }
    }
}

// ---------------------------------------------------------------------------
// MultiGpuFftConfig -- configuration for multi-GPU FFT
// ---------------------------------------------------------------------------

/// Configuration for a multi-GPU FFT operation.
#[derive(Debug, Clone)]
pub struct MultiGpuFftConfig {
    /// Total FFT size (number of complex elements).
    pub total_n: usize,
    /// Number of GPU devices to use.
    pub num_devices: usize,
    /// Floating-point precision.
    pub precision: FftPrecision,
    /// Transform type (C2C, R2C, C2R).
    pub transform_type: FftType,
}

impl MultiGpuFftConfig {
    /// Creates a new multi-GPU FFT configuration.
    pub fn new(
        total_n: usize,
        num_devices: usize,
        precision: FftPrecision,
        transform_type: FftType,
    ) -> Self {
        Self {
            total_n,
            num_devices,
            precision,
            transform_type,
        }
    }

    /// Validates the configuration.
    ///
    /// # Errors
    ///
    /// Returns [`FftError::InvalidSize`] if:
    /// - `total_n` is zero
    /// - `num_devices` is less than 2
    /// - `total_n < num_devices`
    pub fn validate(&self) -> FftResult<()> {
        if self.total_n == 0 {
            return Err(FftError::InvalidSize("total_n must be > 0".to_string()));
        }
        if self.num_devices < 2 {
            return Err(FftError::InvalidSize(
                "multi-GPU FFT requires at least 2 devices".to_string(),
            ));
        }
        if self.total_n < self.num_devices {
            return Err(FftError::InvalidSize(format!(
                "total_n ({}) must be >= num_devices ({})",
                self.total_n, self.num_devices
            )));
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// MultiGpuFftPlan -- execution plan for multi-GPU FFT
// ---------------------------------------------------------------------------

/// A fully specified multi-GPU FFT execution plan.
///
/// The plan holds the decomposition, per-device local FFT plans,
/// and the transpose strategy. Execution follows three phases:
///
/// 1. **Local FFT (Phase 1):** Each device computes FFTs on its slab.
/// 2. **Global transpose:** All-to-all data redistribution.
/// 3. **Local FFT (Phase 2):** Each device computes FFTs on transposed data.
#[derive(Debug)]
pub struct MultiGpuFftPlan {
    /// The validated configuration.
    config: MultiGpuFftConfig,
    /// The slab decomposition across devices.
    decomposition: SlabDecomposition,
    /// One local FFT plan per device.
    local_plans: Vec<FftPlan>,
    /// Strategy for the transpose phase.
    transpose_strategy: TransposeStrategy,
}

impl MultiGpuFftPlan {
    /// Creates a new multi-GPU FFT plan.
    ///
    /// This validates the configuration, creates the slab decomposition,
    /// and builds per-device local FFT plans.
    ///
    /// # Errors
    ///
    /// Returns an error if the configuration is invalid or if any
    /// local plan fails to create.
    pub fn new(config: MultiGpuFftConfig) -> FftResult<Self> {
        config.validate()?;

        let decomposition = SlabDecomposition::new(config.total_n, config.num_devices)?;
        let transpose_strategy = TransposeStrategy::auto_select(config.num_devices);

        let mut local_plans = Vec::with_capacity(config.num_devices);
        for i in 0..config.num_devices {
            let slab_size = decomposition.slab_size(i);
            let plan = FftPlan::new_1d(slab_size, config.transform_type, 1)?
                .with_precision(config.precision);
            local_plans.push(plan);
        }

        Ok(Self {
            config,
            decomposition,
            local_plans,
            transpose_strategy,
        })
    }

    /// Returns the total workspace size in bytes across all devices.
    ///
    /// This includes each device's local FFT workspace plus the
    /// transpose buffer (one full slab per device for staging).
    pub fn total_workspace_bytes(&self) -> usize {
        let local_workspace: usize = self
            .local_plans
            .iter()
            .map(|p| p.estimated_workspace_bytes())
            .sum();

        // Transpose buffer: each device needs space for one full slab
        // to receive transposed data.
        let transpose_buffer: usize = self
            .decomposition
            .slab_sizes
            .iter()
            .map(|&s| s * self.config.precision.complex_bytes())
            .sum();

        local_workspace + transpose_buffer
    }

    /// Returns the local FFT plan for the given device index.
    pub fn local_plan(&self, device_idx: usize) -> Option<&FftPlan> {
        self.local_plans.get(device_idx)
    }

    /// Returns the number of devices in this plan.
    pub fn num_devices(&self) -> usize {
        self.config.num_devices
    }

    /// Returns a reference to the slab decomposition.
    pub fn decomposition(&self) -> &SlabDecomposition {
        &self.decomposition
    }

    /// Returns the transpose strategy.
    pub fn transpose_strategy(&self) -> TransposeStrategy {
        self.transpose_strategy
    }

    /// Returns a human-readable description of the multi-GPU FFT plan.
    pub fn describe(&self) -> String {
        let mut desc = String::new();
        desc.push_str(&format!(
            "MultiGpuFftPlan: N={}, devices={}, precision={:?}, type={}\n",
            self.config.total_n,
            self.config.num_devices,
            self.config.precision,
            self.config.transform_type
        ));
        desc.push_str(&format!("  Decomposition: {}\n", self.decomposition));
        desc.push_str(&format!("  Transpose: {}\n", self.transpose_strategy));
        desc.push_str("  Phases:\n");
        desc.push_str("    1. Local FFT: each device computes FFT on its slab\n");
        desc.push_str("    2. Global transpose: all-to-all redistribution\n");
        desc.push_str("    3. Local FFT: each device computes FFT on transposed slab\n");
        desc.push_str(&format!(
            "  Total workspace: {} bytes\n",
            self.total_workspace_bytes()
        ));

        for (i, plan) in self.local_plans.iter().enumerate() {
            desc.push_str(&format!(
                "  Device {}: slab_size={}, workspace={} bytes\n",
                i,
                self.decomposition.slab_size(i),
                plan.estimated_workspace_bytes()
            ));
        }

        desc
    }
}

// ---------------------------------------------------------------------------
// PhaseDescription -- human-readable description of a phase
// ---------------------------------------------------------------------------

/// Describes one phase of the multi-GPU FFT execution (used for planning
/// and dry-run introspection).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PhaseDescription {
    /// Human-readable name of the phase.
    pub name: String,
    /// Per-device operations in this phase.
    pub device_ops: Vec<String>,
}

// ---------------------------------------------------------------------------
// MultiGpuFftResult -- concrete result of executing the multi-GPU FFT
// ---------------------------------------------------------------------------

/// Timing and throughput information for one phase of the multi-GPU FFT.
#[derive(Debug, Clone)]
pub struct PhaseResult {
    /// Number of elements processed on each device (indexed by device).
    pub elements_per_device: Vec<usize>,
    /// Wall-clock time spent in this phase.
    pub elapsed: Duration,
}

impl PhaseResult {
    /// Returns the total number of elements processed across all devices.
    pub fn total_elements(&self) -> usize {
        self.elements_per_device.iter().sum()
    }
}

/// Result of a complete multi-GPU FFT execution (all three phases).
#[derive(Debug, Clone)]
pub struct MultiGpuFftResult {
    /// Result of Phase 1 (local FFT on slabs).
    pub phase1: PhaseResult,
    /// Result of the global transpose.
    pub transpose: TransposeResult,
    /// Result of Phase 2 (local FFT on transposed data).
    pub phase2: PhaseResult,
    /// Total wall-clock time for all three phases combined.
    pub total_elapsed: Duration,
}

impl MultiGpuFftResult {
    /// Returns the total bytes transferred during the global transpose.
    pub fn total_bytes_transferred(&self) -> usize {
        self.transpose.bytes_transferred
    }
}

/// Result of the global transpose phase.
#[derive(Debug, Clone)]
pub struct TransposeResult {
    /// Descriptions of each device-to-device transfer that was performed.
    pub transfers: Vec<TransferRecord>,
    /// Total bytes moved across device boundaries.
    pub bytes_transferred: usize,
    /// Wall-clock time for the entire transpose.
    pub elapsed: Duration,
}

/// Records one device-to-device transfer in the transpose.
#[derive(Debug, Clone)]
pub struct TransferRecord {
    /// Source device index.
    pub src_device: usize,
    /// Destination device index.
    pub dst_device: usize,
    /// Number of elements transferred.
    pub count: usize,
    /// Byte size of the transfer.
    pub bytes: usize,
}

// ---------------------------------------------------------------------------
// MultiGpuFftExecutor -- orchestrates multi-GPU FFT execution
// ---------------------------------------------------------------------------

/// Orchestrates the execution of a multi-GPU FFT.
///
/// The executor holds a [`MultiGpuFftPlan`] and exposes:
/// - **describe_*** methods for dry-run introspection (no GPU required).
/// - **execute()** for driving real FFT computation across live GPU contexts.
///
/// # Slab FFT Algorithm
///
/// The executor implements the standard three-phase slab FFT:
///
/// 1. **Phase 1** — each device runs a local 1-D FFT along the fast dimension
///    of its slab using [`FftHandle::execute`].
/// 2. **All-to-all transpose** — peer-to-peer (or staged-via-host) copies
///    redistribute data so each device holds a contiguous range of a different
///    dimension.
/// 3. **Phase 2** — each device runs a local 1-D FFT on its newly-received
///    transposed slab.
///
/// # Device Buffers
///
/// Callers supply one `DeviceBuffer<u8>` per device. Buffers must be large
/// enough to hold the device's slab of complex elements in the precision used
/// by the plan. For FP32 complex, each element is 8 bytes; for FP64, 16 bytes.
///
/// A second set of `transpose_bufs` (one per device) is required during Phase
/// 2. These can be the same buffers if execution is synchronised after the
/// transpose.
#[derive(Debug)]
pub struct MultiGpuFftExecutor {
    /// The execution plan.
    plan: MultiGpuFftPlan,
}

impl MultiGpuFftExecutor {
    /// Creates a new executor from a plan.
    pub fn new(plan: MultiGpuFftPlan) -> Self {
        Self { plan }
    }

    /// Returns a reference to the underlying plan.
    pub fn plan(&self) -> &MultiGpuFftPlan {
        &self.plan
    }

    // -----------------------------------------------------------------------
    // Dry-run introspection (no GPU required)
    // -----------------------------------------------------------------------

    /// Describes Phase 1: local FFT on each device's slab (dry run, no GPU).
    pub fn describe_phase1(&self) -> PhaseDescription {
        let mut ops = Vec::with_capacity(self.plan.num_devices());
        for i in 0..self.plan.num_devices() {
            let slab_size = self.plan.decomposition().slab_size(i);
            ops.push(format!(
                "Device {i}: compute local FFT on {slab_size} elements"
            ));
        }
        PhaseDescription {
            name: "Phase 1: Local FFT".to_string(),
            device_ops: ops,
        }
    }

    /// Describes the global transpose phase (dry run, no GPU).
    pub fn describe_transpose(&self) -> PhaseDescription {
        let n = self.plan.num_devices();
        let strategy = self.plan.transpose_strategy();
        let mut ops = Vec::with_capacity(n * n);

        for src in 0..n {
            for dst in 0..n {
                if src == dst {
                    continue;
                }
                if let Some(region) = self.plan.decomposition().transpose_map(src, dst) {
                    ops.push(format!(
                        "Device {src} -> Device {dst}: transfer {} elements \
                         (src_off={}, dst_off={}) via {strategy}",
                        region.count, region.src_offset, region.dst_offset
                    ));
                }
            }
        }

        PhaseDescription {
            name: "Global Transpose".to_string(),
            device_ops: ops,
        }
    }

    /// Describes Phase 2: local FFT on transposed data (dry run, no GPU).
    pub fn describe_phase2(&self) -> PhaseDescription {
        let mut ops = Vec::with_capacity(self.plan.num_devices());
        for i in 0..self.plan.num_devices() {
            let slab_size = self.plan.decomposition().slab_size(i);
            ops.push(format!(
                "Device {i}: compute local FFT on transposed {slab_size} elements"
            ));
        }
        PhaseDescription {
            name: "Phase 2: Local FFT (transposed)".to_string(),
            device_ops: ops,
        }
    }

    // -----------------------------------------------------------------------
    // Real execution
    // -----------------------------------------------------------------------

    /// Executes the full three-phase multi-GPU FFT.
    ///
    /// # Arguments
    ///
    /// * `handles` - One [`FftHandle`] per device (must match `plan.num_devices()`).
    ///   Each handle must be bound to the corresponding device's CUDA context.
    /// * `slab_bufs` - One [`DeviceBuffer<u8>`] per device holding the raw
    ///   byte contents of that device's slab.  Buffers are updated in-place.
    ///   After execution the buffers contain the frequency-domain output.
    /// * `direction` - [`FftDirection::Forward`] or [`FftDirection::Inverse`].
    ///
    /// # Errors
    ///
    /// * [`FftError::InvalidSize`] if `handles.len()` or `slab_bufs.len()` does
    ///   not equal `plan.num_devices()`.
    /// * Any error returned by [`FftHandle::execute`] during phase 1 or phase 2.
    /// * Any CUDA error during the all-to-all transpose (wrapped as
    ///   [`FftError::Cuda`]).
    pub fn execute(
        &self,
        handles: &[FftHandle],
        slab_bufs: &mut [DeviceBuffer<u8>],
        direction: FftDirection,
    ) -> FftResult<MultiGpuFftResult> {
        let n_dev = self.plan.num_devices();

        if handles.len() != n_dev {
            return Err(FftError::InvalidSize(format!(
                "execute: expected {n_dev} FftHandles, got {}",
                handles.len()
            )));
        }
        if slab_bufs.len() != n_dev {
            return Err(FftError::InvalidSize(format!(
                "execute: expected {n_dev} slab buffers, got {}",
                slab_bufs.len()
            )));
        }

        let total_start = Instant::now();

        // ---- Phase 1: local FFT on each slab --------------------------------
        let p1_start = Instant::now();
        let mut p1_elements: Vec<usize> = Vec::with_capacity(n_dev);

        for i in 0..n_dev {
            let slab_size = self.plan.decomposition().slab_size(i);
            let handle = handles
                .get(i)
                .ok_or_else(|| FftError::InternalError(format!("missing handle for device {i}")))?;
            let local_plan = self.plan.local_plan(i).ok_or_else(|| {
                FftError::InternalError(format!("missing local plan for device {i}"))
            })?;
            let buf = slab_bufs.get(i).ok_or_else(|| {
                FftError::InternalError(format!("missing slab buffer for device {i}"))
            })?;

            // The device pointer for the slab: the buffer holds raw bytes so we
            // reinterpret the base pointer as a complex-element pointer for the
            // kernel.  Input and output are the same buffer (in-place FFT).
            let dev_ptr: CUdeviceptr = buf.as_device_ptr();
            handle.execute(local_plan, dev_ptr, dev_ptr, direction)?;

            p1_elements.push(slab_size);
        }

        let phase1 = PhaseResult {
            elements_per_device: p1_elements,
            elapsed: p1_start.elapsed(),
        };

        // ---- All-to-all transpose -------------------------------------------
        let tr_start = Instant::now();
        let mut transfers: Vec<TransferRecord> = Vec::with_capacity(n_dev * n_dev);
        let bytes_per_element = self.plan.config.precision.complex_bytes();
        let strategy = self.plan.transpose_strategy();

        // Collect all (src, dst, region) pairs first to avoid simultaneous
        // mutable and immutable borrows of `slab_bufs`.
        let mut pending: Vec<(usize, usize, TransposeRegion)> =
            Vec::with_capacity(n_dev * (n_dev - 1));
        for src in 0..n_dev {
            for dst in 0..n_dev {
                if src == dst {
                    continue;
                }
                if let Some(region) = self.plan.decomposition().transpose_map(src, dst) {
                    pending.push((src, dst, region));
                }
            }
        }

        for (src, dst, region) in &pending {
            // `TransposeRegion` counts/offsets are in complex elements; the
            // slab buffers are raw `u8`, so convert to byte units here.
            let byte_count = region.count * bytes_per_element;
            let src_byte_offset = region.src_offset * bytes_per_element;
            let dst_byte_offset = region.dst_offset * bytes_per_element;
            transfers.push(TransferRecord {
                src_device: *src,
                dst_device: *dst,
                count: region.count,
                bytes: byte_count,
            });

            // Resolve the real CUDA devices backing the source and
            // destination handles for this transfer pair.  Each handle is
            // bound to its own context, and that context records the device
            // it was created on.
            let src_device = *handles
                .get(*src)
                .ok_or_else(|| FftError::InternalError(format!("missing handle for device {src}")))?
                .context()
                .device();
            let dst_device = *handles
                .get(*dst)
                .ok_or_else(|| FftError::InternalError(format!("missing handle for device {dst}")))?
                .context()
                .device();

            match strategy {
                TransposeStrategy::PeerToPeer => {
                    // Borrow the source (immutable) and destination (mutable)
                    // slab buffers simultaneously via `split_at_mut`.
                    let (lo, hi) = if src < dst {
                        (*src, *dst)
                    } else {
                        (*dst, *src)
                    };
                    let (left, right) = slab_bufs.split_at_mut(lo + 1);
                    let (lo_buf, hi_buf) = (&mut left[lo], &mut right[hi - lo - 1]);
                    let (src_buf, dst_buf) = if *src < *dst {
                        (lo_buf as &DeviceBuffer<u8>, hi_buf)
                    } else {
                        (hi_buf as &DeviceBuffer<u8>, lo_buf)
                    };

                    // Transfer exactly the `TransposeRegion` sub-slice
                    // (`byte_count` bytes at `src_byte_offset` -> `dst_byte_offset`)
                    // between the resolved devices.
                    let copy_result = peer_copy::copy_peer_region(
                        dst_buf,
                        &dst_device,
                        dst_byte_offset,
                        src_buf,
                        &src_device,
                        src_byte_offset,
                        byte_count,
                    );

                    // If direct P2P access is unavailable for this device
                    // pair, fall back to a host-staged transfer rather than
                    // failing the whole FFT.
                    match copy_result {
                        Ok(()) => {}
                        Err(
                            oxicuda_driver::CudaError::NotSupported
                            | oxicuda_driver::CudaError::PeerAccessUnsupported
                            | oxicuda_driver::CudaError::PeerAccessNotEnabled,
                        ) => {
                            let src_handle = handles.get(*src).ok_or_else(|| {
                                FftError::InternalError(format!("missing handle {src}"))
                            })?;
                            let dst_handle = handles.get(*dst).ok_or_else(|| {
                                FftError::InternalError(format!("missing handle {dst}"))
                            })?;
                            stage_via_host(
                                dst_buf,
                                dst_byte_offset,
                                src_buf,
                                src_byte_offset,
                                byte_count,
                                src_handle,
                                dst_handle,
                            )?;
                        }
                        Err(e) => return Err(FftError::Cuda(e)),
                    }
                }
                TransposeStrategy::StagedViaHost => {
                    // Host-staged path: D->H from the source slab into a
                    // pinned buffer, then H->D into the destination slab.
                    let (lo, hi) = if src < dst {
                        (*src, *dst)
                    } else {
                        (*dst, *src)
                    };
                    let (left, right) = slab_bufs.split_at_mut(lo + 1);
                    let (lo_buf, hi_buf) = (&mut left[lo], &mut right[hi - lo - 1]);
                    let (src_buf, dst_buf) = if *src < *dst {
                        (lo_buf as &DeviceBuffer<u8>, hi_buf)
                    } else {
                        (hi_buf as &DeviceBuffer<u8>, lo_buf)
                    };

                    let src_handle = handles
                        .get(*src)
                        .ok_or_else(|| FftError::InternalError(format!("missing handle {src}")))?;
                    let dst_handle = handles
                        .get(*dst)
                        .ok_or_else(|| FftError::InternalError(format!("missing handle {dst}")))?;
                    stage_via_host(
                        dst_buf,
                        dst_byte_offset,
                        src_buf,
                        src_byte_offset,
                        byte_count,
                        src_handle,
                        dst_handle,
                    )?;
                }
            }
        }

        let total_bytes: usize = transfers.iter().map(|t| t.bytes).sum();
        let transpose = TransposeResult {
            transfers,
            bytes_transferred: total_bytes,
            elapsed: tr_start.elapsed(),
        };

        // ---- Phase 2: local FFT on transposed data --------------------------
        let p2_start = Instant::now();
        let mut p2_elements: Vec<usize> = Vec::with_capacity(n_dev);

        for i in 0..n_dev {
            let slab_size = self.plan.decomposition().slab_size(i);
            let handle = handles
                .get(i)
                .ok_or_else(|| FftError::InternalError(format!("missing handle for device {i}")))?;
            let local_plan = self.plan.local_plan(i).ok_or_else(|| {
                FftError::InternalError(format!("missing local plan for device {i}"))
            })?;
            let buf = slab_bufs.get(i).ok_or_else(|| {
                FftError::InternalError(format!("missing slab buffer for device {i}"))
            })?;

            let dev_ptr: CUdeviceptr = buf.as_device_ptr();
            handle.execute(local_plan, dev_ptr, dev_ptr, direction)?;

            p2_elements.push(slab_size);
        }

        let phase2 = PhaseResult {
            elements_per_device: p2_elements,
            elapsed: p2_start.elapsed(),
        };

        Ok(MultiGpuFftResult {
            phase1,
            transpose,
            phase2,
            total_elapsed: total_start.elapsed(),
        })
    }
}

// ---------------------------------------------------------------------------
// Host-staged inter-device transfer
// ---------------------------------------------------------------------------

/// Transfers a contiguous byte sub-range between two device slabs by staging
/// through pinned host memory.
///
/// This is the fallback path used when direct GPU-to-GPU peer access is
/// unavailable (e.g. across NUMA nodes without NVLink), and the primary path
/// for the [`TransposeStrategy::StagedViaHost`] strategy.
///
/// The transfer happens in two legs:
///
/// 1. **Device -> Host:** `byte_count` bytes are copied from `src` starting at
///    `src_byte_offset` into a freshly-allocated [`PinnedBuffer`].  The copy
///    runs on `src_handle`'s stream with `src_handle`'s context current.
/// 2. **Host -> Device:** the staged bytes are copied into `dst` starting at
///    `dst_byte_offset`.  This leg runs on `dst_handle`'s stream with
///    `dst_handle`'s context current.
///
/// Each leg's stream is synchronised before the next step, so the pinned
/// staging buffer is guaranteed to be fully written before it is read and
/// stays alive for the entire transfer.
///
/// # Errors
///
/// Returns [`FftError::Cuda`] if pinned allocation, either copy, a context
/// switch, or stream synchronisation fails.
fn stage_via_host(
    dst: &mut DeviceBuffer<u8>,
    dst_byte_offset: usize,
    src: &DeviceBuffer<u8>,
    src_byte_offset: usize,
    byte_count: usize,
    src_handle: &FftHandle,
    dst_handle: &FftHandle,
) -> FftResult<()> {
    // A zero-byte region is a well-defined no-op.
    if byte_count == 0 {
        return Ok(());
    }

    // Pinned host staging buffer sized exactly to the transfer.
    let mut staging: PinnedBuffer<u8> = PinnedBuffer::alloc(byte_count).map_err(FftError::Cuda)?;

    // ---- Leg 1: Device -> Host on the source device ----------------------
    src_handle.context().set_current().map_err(FftError::Cuda)?;
    copy::copy_dtoh_region_async(
        &mut staging,
        src,
        src_byte_offset,
        byte_count,
        src_handle.stream(),
    )
    .map_err(FftError::Cuda)?;
    // The staging buffer must be fully populated before the H->D leg reads it.
    src_handle.stream().synchronize().map_err(FftError::Cuda)?;

    // ---- Leg 2: Host -> Device on the destination device -----------------
    dst_handle.context().set_current().map_err(FftError::Cuda)?;
    copy::copy_htod_region_async(
        dst,
        dst_byte_offset,
        &staging,
        byte_count,
        dst_handle.stream(),
    )
    .map_err(FftError::Cuda)?;
    // Synchronise so the staging buffer is not freed (on drop) while the
    // asynchronous H->D copy is still in flight.
    dst_handle.stream().synchronize().map_err(FftError::Cuda)?;

    Ok(())
}

// ---------------------------------------------------------------------------
// NvLink Overlap Pipeline
// ---------------------------------------------------------------------------

/// Pipeline stage descriptor for compute/transfer overlap.
#[derive(Debug, Clone)]
pub struct OverlapStage {
    /// Human-readable stage name.
    pub name: String,
    /// Device index responsible for this stage.
    pub device_idx: usize,
    /// Number of elements processed in this stage.
    pub elements: usize,
    /// Whether compute and transfer happen simultaneously in this stage.
    pub overlapped: bool,
}

/// A double-buffered NVLink pipeline plan for overlapping FFT compute with
/// data transfers.
///
/// For an N-point FFT distributed across P devices over NVLink, this plan
/// divides the slab into two halves per device (ping and pong buffers).
/// While device `i` computes its FFT on the "ping" half, it simultaneously
/// receives the next "pong" data from `i-1` over NVLink.
///
/// This achieves near-linear scaling on NVLink-connected GPU systems by
/// hiding the transfer latency.
///
/// # Double-buffer layout (per device)
///
/// ```text
/// time →  [FFT ping_0] [FFT ping_1] [FFT ping_2] …
///         [  send pong_0  ][  send pong_1  ] …
/// ```
#[derive(Debug)]
pub struct NvLinkOverlapPlan {
    /// Number of GPU devices in the NVLink fabric.
    pub num_devices: usize,
    /// Elements per device per half-buffer (the "ping" or "pong" slice).
    pub half_slab_elements: usize,
    /// Number of pipeline stages (= `num_devices * 2` for full overlap).
    pub num_stages: usize,
    /// Ordered list of pipeline stages.
    pub stages: Vec<OverlapStage>,
    /// Whether the active hardware topology supports full NVLink overlap.
    /// When `false`, the pipeline degrades to sequential PeerToPeer.
    pub nvlink_available: bool,
}

impl NvLinkOverlapPlan {
    /// Build a double-buffered NVLink overlap plan for `num_devices` GPUs
    /// and a total FFT of `total_n` elements.
    ///
    /// # Errors
    ///
    /// Returns [`FftError::InvalidSize`] if `total_n < num_devices * 2`
    /// (each buffer half must have at least 1 element per device).
    pub fn new(num_devices: usize, total_n: usize, nvlink_available: bool) -> FftResult<Self> {
        if num_devices < 2 {
            return Err(FftError::InvalidSize(
                "NvLink overlap pipeline requires at least 2 devices".into(),
            ));
        }
        if total_n < num_devices * 2 {
            return Err(FftError::InvalidSize(format!(
                "total_n ({total_n}) must be >= num_devices * 2 ({})",
                num_devices * 2
            )));
        }

        let slab_n = total_n / num_devices;
        let half_slab_elements = slab_n / 2;

        // Build ordered stages: for each device, we interleave:
        //   Stage 2k  : compute FFT on "ping" half
        //   Stage 2k+1: transfer "pong" half to next device over NVLink
        let num_stages = num_devices * 2;
        let mut stages = Vec::with_capacity(num_stages);

        for dev in 0..num_devices {
            stages.push(OverlapStage {
                name: format!("dev{dev}: FFT ping"),
                device_idx: dev,
                elements: half_slab_elements,
                overlapped: nvlink_available,
            });
            stages.push(OverlapStage {
                name: format!(
                    "dev{dev}: NVLink send pong → dev{}",
                    (dev + 1) % num_devices
                ),
                device_idx: dev,
                elements: half_slab_elements,
                overlapped: nvlink_available,
            });
        }

        Ok(Self {
            num_devices,
            half_slab_elements,
            num_stages,
            stages,
            nvlink_available,
        })
    }

    /// Total compute elements across all pipeline stages.
    pub fn total_elements(&self) -> usize {
        self.stages.iter().map(|s| s.elements).sum()
    }

    /// Return `true` if the pipeline is running in full-overlap mode.
    pub fn is_overlapped(&self) -> bool {
        self.nvlink_available && self.stages.iter().all(|s| s.overlapped)
    }

    /// Estimated speedup factor relative to sequential execution.
    ///
    /// With full NVLink overlap, ideally the transfer is fully hidden under
    /// compute, yielding speedup ≈ 2× relative to a non-pipelined P2P path.
    /// Without NVLink the speedup is 1× (sequential).
    pub fn estimated_speedup(&self) -> f32 {
        if self.nvlink_available { 1.8 } else { 1.0 }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- SlabDecomposition tests --

    #[test]
    fn slab_even_division() {
        let decomp = SlabDecomposition::new(1024, 4);
        assert!(decomp.is_ok());
        let decomp = decomp.ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            assert_eq!(d.slab_sizes, vec![256, 256, 256, 256]);
            assert_eq!(d.slab_offsets, vec![0, 256, 512, 768]);
        }
    }

    #[test]
    fn slab_uneven_division() {
        let decomp = SlabDecomposition::new(10, 3);
        assert!(decomp.is_ok());
        let decomp = decomp.ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            // 10 / 3 = 3 remainder 1 => first device gets 4, rest get 3
            assert_eq!(d.slab_sizes, vec![4, 3, 3]);
            assert_eq!(d.slab_offsets, vec![0, 4, 7]);
            let total: usize = d.slab_sizes.iter().sum();
            assert_eq!(total, 10);
        }
    }

    #[test]
    fn slab_single_device() {
        let decomp = SlabDecomposition::new(512, 1);
        assert!(decomp.is_ok());
        let decomp = decomp.ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            assert_eq!(d.slab_sizes, vec![512]);
            assert_eq!(d.slab_offsets, vec![0]);
        }
    }

    #[test]
    fn slab_offset_computation() {
        let decomp = SlabDecomposition::new(100, 4);
        assert!(decomp.is_ok());
        if let Ok(d) = decomp {
            // 100 / 4 = 25 even
            for i in 0..4 {
                assert_eq!(d.slab_offset(i), i * 25);
                assert_eq!(d.slab_size(i), 25);
            }
            // Out of range
            assert_eq!(d.slab_offset(4), 0);
            assert_eq!(d.slab_size(4), 0);
        }
    }

    #[test]
    fn slab_zero_total_n() {
        let result = SlabDecomposition::new(0, 4);
        assert!(matches!(result, Err(FftError::InvalidSize(_))));
    }

    #[test]
    fn slab_zero_devices() {
        let result = SlabDecomposition::new(100, 0);
        assert!(matches!(result, Err(FftError::InvalidSize(_))));
    }

    #[test]
    fn slab_n_less_than_devices() {
        let result = SlabDecomposition::new(3, 8);
        assert!(matches!(result, Err(FftError::InvalidSize(_))));
    }

    // -- TransposeMap tests --

    #[test]
    fn transpose_map_same_device() {
        let decomp = SlabDecomposition::new(1024, 4).ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            let region = d.transpose_map(0, 0);
            assert!(region.is_some());
            if let Some(r) = region {
                // Self-transfer: the intersection is the entire slab
                assert_eq!(r.src_offset, 0);
                assert_eq!(r.dst_offset, 0);
                assert_eq!(r.count, 256);
            }
        }
    }

    #[test]
    fn transpose_map_adjacent_devices() {
        let decomp = SlabDecomposition::new(1024, 4).ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            // Device 0 owns [0, 256), device 1 owns [256, 512)
            // No overlap -> falls into compute_transpose_count path
            let region = d.transpose_map(0, 1);
            assert!(region.is_some());
            if let Some(r) = region {
                assert!(r.count > 0);
            }
        }
    }

    #[test]
    fn transpose_map_out_of_range() {
        let decomp = SlabDecomposition::new(1024, 4).ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            assert!(d.transpose_map(0, 10).is_none());
            assert!(d.transpose_map(10, 0).is_none());
        }
    }

    // -- Config validation tests --

    #[test]
    fn config_validate_zero_n() {
        let config = MultiGpuFftConfig::new(0, 4, FftPrecision::Single, FftType::C2C);
        assert!(config.validate().is_err());
    }

    #[test]
    fn config_validate_one_device() {
        let config = MultiGpuFftConfig::new(1024, 1, FftPrecision::Single, FftType::C2C);
        assert!(config.validate().is_err());
    }

    #[test]
    fn config_validate_n_less_than_devices() {
        let config = MultiGpuFftConfig::new(3, 8, FftPrecision::Single, FftType::C2C);
        assert!(config.validate().is_err());
    }

    #[test]
    fn config_validate_valid() {
        let config = MultiGpuFftConfig::new(1024, 4, FftPrecision::Single, FftType::C2C);
        assert!(config.validate().is_ok());
    }

    // -- MultiGpuFftPlan tests --

    #[test]
    fn plan_creation() {
        let config = MultiGpuFftConfig::new(1024, 4, FftPrecision::Single, FftType::C2C);
        let plan = MultiGpuFftPlan::new(config);
        assert!(plan.is_ok());
        if let Ok(p) = plan {
            assert_eq!(p.num_devices(), 4);
            assert!(p.local_plan(0).is_some());
            assert!(p.local_plan(3).is_some());
            assert!(p.local_plan(4).is_none());
        }
    }

    #[test]
    fn plan_workspace_calculation() {
        let config = MultiGpuFftConfig::new(1024, 4, FftPrecision::Single, FftType::C2C);
        let plan = MultiGpuFftPlan::new(config);
        assert!(plan.is_ok());
        if let Ok(p) = plan {
            let workspace = p.total_workspace_bytes();
            // At minimum, transpose buffers: 1024 * 8 bytes (f32 complex)
            assert!(workspace > 0);
        }
    }

    #[test]
    fn plan_describe_output() {
        let config = MultiGpuFftConfig::new(1024, 2, FftPrecision::Single, FftType::C2C);
        let plan = MultiGpuFftPlan::new(config);
        assert!(plan.is_ok());
        if let Ok(p) = plan {
            let desc = p.describe();
            assert!(desc.contains("MultiGpuFftPlan"));
            assert!(desc.contains("N=1024"));
            assert!(desc.contains("devices=2"));
            assert!(desc.contains("1. Local FFT"));
            assert!(desc.contains("2. Global transpose"));
            assert!(desc.contains("transpose"));
        }
    }

    #[test]
    fn plan_double_precision() {
        let config = MultiGpuFftConfig::new(512, 2, FftPrecision::Double, FftType::C2C);
        let plan = MultiGpuFftPlan::new(config);
        assert!(plan.is_ok());
        if let Ok(p) = plan {
            let workspace = p.total_workspace_bytes();
            // Double precision: 16 bytes per complex element, 512 elements
            assert!(workspace >= 512 * 16);
        }
    }

    // -- TransposeStrategy tests --

    #[test]
    fn transpose_strategy_auto_select() {
        assert_eq!(
            TransposeStrategy::auto_select(1),
            TransposeStrategy::PeerToPeer
        );
        assert_eq!(
            TransposeStrategy::auto_select(2),
            TransposeStrategy::PeerToPeer
        );
        assert_eq!(
            TransposeStrategy::auto_select(3),
            TransposeStrategy::StagedViaHost
        );
        assert_eq!(
            TransposeStrategy::auto_select(8),
            TransposeStrategy::StagedViaHost
        );
    }

    // -- MultiGpuFftExecutor tests --

    #[test]
    fn executor_phase_descriptions() {
        let config = MultiGpuFftConfig::new(1024, 2, FftPrecision::Single, FftType::C2C);
        let plan = MultiGpuFftPlan::new(config);
        assert!(plan.is_ok());
        if let Ok(p) = plan {
            let executor = MultiGpuFftExecutor::new(p);

            let p1 = executor.describe_phase1();
            assert_eq!(p1.name, "Phase 1: Local FFT");
            assert_eq!(p1.device_ops.len(), 2);

            let transpose = executor.describe_transpose();
            assert_eq!(transpose.name, "Global Transpose");
            // 2 devices, each sends to the other = 2 transfers
            assert_eq!(transpose.device_ops.len(), 2);

            let p2 = executor.describe_phase2();
            assert_eq!(p2.name, "Phase 2: Local FFT (transposed)");
            assert_eq!(p2.device_ops.len(), 2);
        }
    }

    /// Verify that N=1024 with P=2 devices decomposes into exactly 2 slabs of 512.
    #[test]
    fn plan_1024_2devices_correct_decomposition() {
        let config = MultiGpuFftConfig::new(1024, 2, FftPrecision::Single, FftType::C2C);
        let plan = MultiGpuFftPlan::new(config).expect("plan creation must succeed");

        let decomp = plan.decomposition();
        assert_eq!(decomp.num_devices, 2);
        assert_eq!(decomp.total_n, 1024);
        assert_eq!(decomp.slab_sizes, vec![512, 512]);
        assert_eq!(decomp.slab_offsets, vec![0, 512]);

        // Phase descriptions confirm each device operates on 512 elements.
        let executor = MultiGpuFftExecutor::new(plan);
        let p1 = executor.describe_phase1();
        assert_eq!(p1.device_ops.len(), 2);
        assert!(p1.device_ops[0].contains("512"));
        assert!(p1.device_ops[1].contains("512"));
    }

    /// Verify that the transpose region count is symmetric:
    /// device A sends as many elements to device B as B sends to A.
    #[test]
    fn transpose_region_symmetry() {
        // Use an even decomposition so symmetry is exact.
        let decomp = SlabDecomposition::new(1024, 4).expect("decomp must succeed");

        for src in 0..decomp.num_devices {
            for dst in 0..decomp.num_devices {
                if src == dst {
                    continue;
                }
                let r_forward = decomp
                    .transpose_map(src, dst)
                    .expect("transpose_map must return Some for valid indices");
                let r_reverse = decomp
                    .transpose_map(dst, src)
                    .expect("transpose_map must return Some for valid indices");

                assert_eq!(
                    r_forward.count, r_reverse.count,
                    "transfer count should be symmetric: ({src}->{dst})={} vs ({dst}->{src})={}",
                    r_forward.count, r_reverse.count
                );
            }
        }
    }

    /// Verify that `PhaseResult::total_elements()` sums correctly.
    #[test]
    fn phase_result_total_elements() {
        let result = PhaseResult {
            elements_per_device: vec![512, 512],
            elapsed: std::time::Duration::ZERO,
        };
        assert_eq!(result.total_elements(), 1024);
    }

    /// Verify `MultiGpuFftResult::total_bytes_transferred()` delegates to transpose.
    #[test]
    fn multi_gpu_result_bytes_transferred() {
        let result = MultiGpuFftResult {
            phase1: PhaseResult {
                elements_per_device: vec![512],
                elapsed: std::time::Duration::ZERO,
            },
            transpose: TransposeResult {
                transfers: vec![TransferRecord {
                    src_device: 0,
                    dst_device: 1,
                    count: 256,
                    bytes: 2048,
                }],
                bytes_transferred: 2048,
                elapsed: std::time::Duration::ZERO,
            },
            phase2: PhaseResult {
                elements_per_device: vec![512],
                elapsed: std::time::Duration::ZERO,
            },
            total_elapsed: std::time::Duration::ZERO,
        };
        assert_eq!(result.total_bytes_transferred(), 2048);
    }

    #[test]
    fn slab_display() {
        let decomp = SlabDecomposition::new(1024, 4).ok();
        assert!(decomp.is_some());
        if let Some(d) = decomp {
            let s = format!("{d}");
            assert!(s.contains("1024"));
            assert!(s.contains("4"));
        }
    }

    // -- NvLink overlap plan tests --

    #[test]
    fn nvlink_plan_creates_correct_stages() {
        let plan = NvLinkOverlapPlan::new(4, 1024, true)
            .expect("NvLink overlap plan should succeed for valid device count and size");
        assert_eq!(plan.num_devices, 4);
        assert_eq!(plan.num_stages, 8);
        assert_eq!(plan.half_slab_elements, 128); // (1024/4)/2
        assert_eq!(plan.stages.len(), 8);
    }

    #[test]
    fn nvlink_plan_stage_names() {
        let plan = NvLinkOverlapPlan::new(2, 256, true)
            .expect("NvLink overlap plan should succeed for valid device count and size");
        assert!(plan.stages[0].name.contains("FFT ping"));
        assert!(plan.stages[1].name.contains("NVLink send pong"));
        assert!(plan.stages[2].name.contains("FFT ping"));
    }

    #[test]
    fn nvlink_plan_total_elements() {
        let plan = NvLinkOverlapPlan::new(2, 256, true)
            .expect("NvLink overlap plan should succeed for valid device count and size");
        // 4 stages × 64 elements each = 256
        assert_eq!(plan.total_elements(), 256);
    }

    #[test]
    fn nvlink_plan_is_overlapped_when_available() {
        let plan = NvLinkOverlapPlan::new(2, 256, true)
            .expect("NvLink overlap plan should succeed for valid device count and size");
        assert!(plan.is_overlapped());
        assert!(plan.estimated_speedup() > 1.0);
    }

    #[test]
    fn nvlink_plan_not_overlapped_without_hw() {
        let plan = NvLinkOverlapPlan::new(2, 256, false)
            .expect("NvLink overlap plan should succeed for valid device count and size");
        assert!(!plan.is_overlapped());
        assert_eq!(plan.estimated_speedup(), 1.0);
    }

    #[test]
    fn nvlink_plan_rejects_single_device() {
        let result = NvLinkOverlapPlan::new(1, 256, true);
        assert!(result.is_err());
    }

    #[test]
    fn nvlink_plan_rejects_too_small_n() {
        // Need total_n >= num_devices * 2 = 6
        let result = NvLinkOverlapPlan::new(3, 4, true);
        assert!(result.is_err());
    }

    // -- Host-staged / peer transpose transfer tests --
    //
    // These exercise the real transpose data movement.  On a single-GPU box
    // the source and destination handles are both bound to device 0, so the
    // transfer becomes a within-device sub-buffer copy — which still proves
    // that exactly `region.count` elements move between the right offsets.

    /// Builds an [`FftHandle`] on device 0, or returns `None` so the caller
    /// skips the test when no CUDA device is available.
    #[cfg(test)]
    fn gpu_handle() -> Option<FftHandle> {
        use std::sync::Arc;
        if oxicuda_driver::init().is_err() {
            return None;
        }
        let device = oxicuda_driver::device::Device::get(0).ok()?;
        let context = Arc::new(oxicuda_driver::Context::new(&device).ok()?);
        FftHandle::new(&context).ok()
    }

    #[test]
    fn stage_via_host_moves_exact_region() {
        let Some(src_handle) = gpu_handle() else {
            eprintln!("skipping: no CUDA device");
            return;
        };
        let Some(dst_handle) = gpu_handle() else {
            eprintln!("skipping: no CUDA device");
            return;
        };

        // Source slab: bytes 0,1,2,...,15.  Destination slab: all zero.
        let host_src: Vec<u8> = (0..16u8).collect();
        let Ok(src_buf) = DeviceBuffer::<u8>::from_host(&host_src) else {
            eprintln!("skipping: device alloc failed");
            return;
        };
        let Ok(mut dst_buf) = DeviceBuffer::<u8>::from_host(&[0u8; 16]) else {
            eprintln!("skipping: device alloc failed");
            return;
        };

        // Move 4 bytes from src[4..8] -> dst[10..14] via host staging.
        if stage_via_host(&mut dst_buf, 10, &src_buf, 4, 4, &src_handle, &dst_handle).is_err() {
            eprintln!("skipping: host-staged transfer failed");
            return;
        }

        let mut out = [0u8; 16];
        if dst_buf.copy_to_host(&mut out).is_err() {
            eprintln!("skipping: copy back failed");
            return;
        }
        // Exactly dst[10..14] must hold src[4..8] = [4,5,6,7]; nothing else.
        let mut expected = [0u8; 16];
        expected[10..14].copy_from_slice(&[4, 5, 6, 7]);
        assert_eq!(out, expected);
    }

    #[test]
    fn stage_via_host_zero_count_is_noop() {
        let Some(src_handle) = gpu_handle() else {
            eprintln!("skipping: no CUDA device");
            return;
        };
        let Some(dst_handle) = gpu_handle() else {
            eprintln!("skipping: no CUDA device");
            return;
        };
        let Ok(src_buf) = DeviceBuffer::<u8>::from_host(&[1u8, 2, 3, 4]) else {
            eprintln!("skipping: device alloc failed");
            return;
        };
        let Ok(mut dst_buf) = DeviceBuffer::<u8>::from_host(&[9u8, 9, 9, 9]) else {
            eprintln!("skipping: device alloc failed");
            return;
        };
        // A zero-byte transfer must leave the destination untouched.
        assert!(stage_via_host(&mut dst_buf, 0, &src_buf, 0, 0, &src_handle, &dst_handle).is_ok());
        let mut out = [0u8; 4];
        if dst_buf.copy_to_host(&mut out).is_err() {
            eprintln!("skipping: copy back failed");
            return;
        }
        assert_eq!(out, [9, 9, 9, 9]);
    }

    #[test]
    fn executor_transpose_moves_data_single_gpu() {
        // End-to-end check that the transpose phase of `execute` actually
        // moves bytes between slabs.  We use a tiny 2-device plan and a
        // forward FFT; both handles are bound to device 0.
        let Some(handle_a) = gpu_handle() else {
            eprintln!("skipping: no CUDA device");
            return;
        };
        let Some(handle_b) = gpu_handle() else {
            eprintln!("skipping: no CUDA device");
            return;
        };

        let config = MultiGpuFftConfig::new(8, 2, FftPrecision::Single, FftType::C2C);
        let Ok(plan) = MultiGpuFftPlan::new(config) else {
            eprintln!("skipping: plan creation failed");
            return;
        };
        let executor = MultiGpuFftExecutor::new(plan);

        // Each slab holds 4 complex f32 elements = 32 bytes.
        let slab_bytes = 4 * FftPrecision::Single.complex_bytes();
        let Ok(buf0) = DeviceBuffer::<u8>::zeroed(slab_bytes) else {
            eprintln!("skipping: device alloc failed");
            return;
        };
        let Ok(buf1) = DeviceBuffer::<u8>::zeroed(slab_bytes) else {
            eprintln!("skipping: device alloc failed");
            return;
        };
        let mut slabs = [buf0, buf1];
        let handles = [handle_a, handle_b];

        match executor.execute(&handles, &mut slabs, FftDirection::Forward) {
            Ok(result) => {
                // The transpose must report the byte total it actually moved.
                let recorded: usize = result.transpose.transfers.iter().map(|t| t.bytes).sum();
                assert_eq!(recorded, result.transpose.bytes_transferred);
                // With 2 devices each sends one region to the other.
                assert_eq!(result.transpose.transfers.len(), 2);
            }
            Err(e) => {
                // Kernel JIT / launch may be unavailable in some CI images.
                eprintln!("skipping: multi-GPU FFT execution failed: {e}");
            }
        }
    }
}
