//! Multi-GPU support for data-parallel training.
//!
//! Provides device enumeration, configuration, and gradient synchronization
//! for single-node multi-GPU training (e.g., 2-8 GPUs with NVLink).
//!
//! # Feature gates
//!
//! - **Default** (no feature): `MultiGpuConfig::detect()` probes for multiple
//!   GPUs and returns `None` on single-GPU setups.
//! - **`nccl`**: Enables `NcclGradientSync` for all-reduce gradient averaging
//!   across devices.  Requires the NCCL library installed on the system.

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

use crate::MLError;

/// Configuration for multi-GPU data-parallel training.
#[derive(Debug, Clone)]
pub struct MultiGpuConfig {
    /// Available CUDA streams (one per GPU)
    pub streams: Vec<Arc<CudaStream>>,
    /// Synchronize gradients every N optimizer steps (default: 1)
    pub sync_every_n_steps: usize,
    /// World size (number of GPUs)
    pub world_size: usize,
}

impl MultiGpuConfig {
    /// Detect available GPUs and return a config if more than one is found.
    ///
    /// Returns `Ok(None)` on single-GPU or CPU-only setups.
    pub fn detect() -> Result<Option<Self>, MLError> {
        let gpu_count = Self::count_cuda_devices();
        if gpu_count <= 1 {
            return Ok(None);
        }

        let mut streams = Vec::with_capacity(gpu_count);
        for i in 0..gpu_count {
            let context = CudaContext::new(i).map_err(|e| {
                MLError::ModelError(format!("Failed to init CUDA device {i}: {e}"))
            })?;
            streams.push(context.default_stream());
        }

        info!(
            "Multi-GPU detected: {} devices available for data parallelism",
            streams.len()
        );

        Ok(Some(Self {
            world_size: streams.len(),
            streams,
            sync_every_n_steps: 1,
        }))
    }

    /// Count available CUDA devices (probes ordinals 0..8). Public for hyperopt multi-GPU.
    pub fn count_cuda_devices_pub() -> usize {
        Self::count_cuda_devices()
    }

    /// Count available CUDA devices (probes ordinals 0..8).
    fn count_cuda_devices() -> usize {
        // Use cuDeviceGetCount — don't probe with CudaContext::new()
        // because failed probes (e.g. device 1 on single-GPU) pollute
        // the cudarc error_state and invalidate CUDA Graph capture.
        CudaContext::device_count().unwrap_or(0) as usize
    }

    /// Shard a dataset across devices by splitting into equal-sized chunks.
    ///
    /// Returns a Vec of index ranges, one per device.
    pub fn shard_indices(&self, total_samples: usize) -> Vec<std::ops::Range<usize>> {
        let chunk = total_samples / self.world_size;
        let remainder = total_samples % self.world_size;
        let mut ranges = Vec::with_capacity(self.world_size);
        let mut start = 0;
        for i in 0..self.world_size {
            let extra = if i < remainder { 1 } else { 0 };
            let end = start + chunk + extra;
            ranges.push(start..end);
            start = end;
        }
        ranges
    }
}

// ---------------------------------------------------------------------------
// NCCL gradient synchronization (requires `nccl` feature + NCCL library)
// ---------------------------------------------------------------------------

/// Placeholder for NCCL-backed gradient synchronization.
///
/// The actual implementation requires `cudarc::nccl::Comm` which depends
/// on the NCCL library being installed on the system.  This is gated
/// behind the `nccl` cargo feature.
///
/// # Usage (when feature is enabled)
///
/// ```ignore
/// let sync = NcclGradientSync::new(&multi_gpu_config.devices)?;
/// // After backward pass on each device:
/// sync.all_reduce_gradients(&mut per_device_grads)?;
/// ```
#[cfg(feature = "nccl")]
#[derive(Debug)]
pub struct NcclGradientSync {
    /// Number of GPUs participating in the sync.
    world_size: usize,
}

#[cfg(feature = "nccl")]
impl NcclGradientSync {
    /// Initialize NCCL communicators for all devices.
    ///
    /// # Errors
    ///
    /// Returns `MLError::ModelError` if NCCL initialization fails
    /// (e.g., NCCL library not found or GPU topology incompatible).
    pub fn new(devices: &[Arc<CudaStream>]) -> Result<Self, MLError> {
        let world_size = devices.len();
        if world_size < 2 {
            return Err(MLError::ModelError(
                "NCCL requires at least 2 devices".into(),
            ));
        }
        info!("NCCL gradient sync initialized for {} devices", world_size);
        Ok(Self { world_size })
    }

    /// World size (number of GPUs).
    pub fn world_size(&self) -> usize {
        self.world_size
    }
}

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

    #[test]
    fn test_multi_gpu_detect_returns_none_on_cpu() {
        // On a CI machine without multiple GPUs, detect() should return None
        let result = MultiGpuConfig::detect();
        assert!(result.is_ok());
        // We can't assert None (might have GPUs), but we can assert it doesn't crash
    }

    #[test]
    fn test_shard_indices_even() {
        let config = MultiGpuConfig {
            streams: vec![],
            sync_every_n_steps: 1,
            world_size: 4,
        };
        let ranges = config.shard_indices(100);
        assert_eq!(ranges.len(), 4);
        assert_eq!(ranges[0], 0..25);
        assert_eq!(ranges[1], 25..50);
        assert_eq!(ranges[2], 50..75);
        assert_eq!(ranges[3], 75..100);
    }

    #[test]
    fn test_shard_indices_uneven() {
        let config = MultiGpuConfig {
            streams: vec![],
            sync_every_n_steps: 1,
            world_size: 3,
        };
        let ranges = config.shard_indices(10);
        assert_eq!(ranges.len(), 3);
        // 10 / 3 = 3 remainder 1 → first device gets 4, others get 3
        assert_eq!(ranges[0], 0..4);
        assert_eq!(ranges[1], 4..7);
        assert_eq!(ranges[2], 7..10);
    }

    #[test]
    fn test_shard_indices_single_device() {
        let config = MultiGpuConfig {
            streams: vec![],
            sync_every_n_steps: 1,
            world_size: 1,
        };
        let ranges = config.shard_indices(50);
        assert_eq!(ranges.len(), 1);
        assert_eq!(ranges[0], 0..50);
    }

    #[test]
    fn test_count_cuda_devices() {
        // Should not crash regardless of GPU availability
        let count = MultiGpuConfig::count_cuda_devices();
        assert!(count <= 8);
    }
}
