//! Comprehensive tests for NCCL backend implementation
//!
//! This module contains tests that validate the NCCL backend functionality,
//! including communicator management, collective operations, and error handling.
#![cfg(feature = "disabled_tests")]

use crate::backend::{Backend, BackendConfig};
use crate::comm::ReduceOp;
use crate::error::Result;
use crate::group::{ProcessGroup, WorldInfo};
use crate::nccl::{NcclBackend, NcclConfig};
use crate::{Device, Tensor, TensorShape};
use proptest::prelude::*;

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

    /// Test NCCL configuration creation and validation
    #[test]
    fn test_nccl_config() {
        let config = NcclConfig::default();
        assert_eq!(config.socket_ifname, "^lo");
        assert_eq!(config.debug_level, "INFO");
        assert_eq!(config.tree_threshold, 0);
        assert_eq!(config.p2p_threshold, 8192);
        assert_eq!(config.buffer_size, 1024 * 1024);
    }

    /// Test NCCL backend configuration creation
    #[test]
    fn test_nccl_backend_config() {
        let config = BackendConfig::nccl();
        assert_eq!(config.backend, Backend::Nccl);
        assert!(config.nccl_config().is_some());

        let nccl_config = config.nccl_config().unwrap();
        assert_eq!(nccl_config.socket_ifname, "^lo");
    }

    /// Test NCCL backend creation
    #[test]
    fn test_nccl_backend_creation() {
        let nccl_config = NcclConfig::default();
        let backend = NcclBackend::new(nccl_config);

        assert!(backend.world_comm().read().is_none());
    }

    /// Test process group creation with NCCL backend (CPU fallback mode)
    #[tokio::test]
    async fn test_process_group_nccl_fallback() {
        let world_info = WorldInfo::new(2, 0, Backend::Nccl);

        // This should work even without CUDA as it will fall back to CPU simulation
        let result = ProcessGroup::new(Backend::Cpu, world_info);

        match result {
            Ok(pg) => {
                assert_eq!(pg.world_size(), 2);
                assert_eq!(pg.rank(), 0);

                // Cleanup
                let _ = pg.cleanup().await;
            }
            Err(_) => {
                // Expected if CUDA/NCCL is not available
                println!("NCCL not available, test skipped");
            }
        }
    }

    /// Test tensor operations with various data types and sizes
    #[tokio::test]
    async fn test_tensor_operations() {
        let world_info = WorldInfo::new(1, 0, Backend::Cpu);
        let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();

        // Test various tensor shapes
        let test_shapes = vec![
            vec![1],
            vec![10],
            vec![128],
            vec![1024],
            vec![2, 3],
            vec![4, 4, 4],
            vec![8, 8, 8, 8],
        ];

        for shape_dims in test_shapes {
            let tensor = Tensor::ones(&shape_dims, &crate::Device::default()).unwrap();
            let mut test_tensor = tensor.clone();

            // Test AllReduce
            let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
            assert!(
                result.is_ok(),
                "AllReduce failed for shape {:?}",
                shape_dims
            );

            // Test Broadcast
            let mut broadcast_tensor = tensor.clone();
            let result = pg.broadcast(&mut broadcast_tensor, 0).await;
            assert!(
                result.is_ok(),
                "Broadcast failed for shape {:?}",
                shape_dims
            );
        }

        pg.cleanup().await.unwrap();
    }

    /// Test error conditions and edge cases
    #[tokio::test]
    async fn test_error_conditions() {
        let world_info = WorldInfo::new(2, 0, Backend::Cpu);
        let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();

        // Test with empty tensor
        let empty_shape = TensorShape::new(vec![0]).unwrap();
        let empty_tensor = Tensor::zeros(empty_shape, &crate::Device::default()).unwrap();
        let mut test_tensor = empty_tensor.clone();

        let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
        // This might succeed or fail depending on implementation - both are acceptable

        // Test broadcast with invalid root
        let valid_tensor = Tensor::ones(&vec![3], &crate::Device::default()).unwrap();
        let mut broadcast_tensor = valid_tensor.clone();
        let result = pg.broadcast(&mut broadcast_tensor, 99).await; // Invalid root
        assert!(result.is_err(), "Broadcast should fail with invalid root");

        pg.cleanup().await.unwrap();
    }

    /// Test process group cleanup
    #[tokio::test]
    async fn test_process_group_cleanup() {
        let world_info = WorldInfo::new(1, 0, Backend::Cpu);
        let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();

        // Test cleanup
        assert!(pg.cleanup().await.is_ok());

        // Test double cleanup (should not error)
        assert!(pg.cleanup().await.is_ok());
    }

    /// Test concurrent operations
    #[tokio::test]
    async fn test_concurrent_operations() {
        let world_info = WorldInfo::new(1, 0, Backend::Cpu);
        let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();

        // Create multiple tasks performing operations concurrently
        let mut tasks = Vec::new();

        for i in 0..5 {
            let pg_clone = pg.clone();
            let task = tokio::spawn(async move {
                let tensor = Tensor::ones(&vec![10], &crate::Device::default()).unwrap();
                let mut test_tensor = tensor.clone();

                // Perform operation
                let result = pg_clone.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
                assert!(result.is_ok(), "Concurrent operation {} failed", i);
            });
            tasks.push(task);
        }

        // Wait for all tasks to complete
        for task in tasks {
            task.await.unwrap();
        }

        pg.cleanup().await.unwrap();
    }

    /// Property-based test for tensor operations
    proptest! {
        #[test]
        fn prop_test_tensor_shapes(
            dims in prop::collection::vec(1usize..100, 1..4),
        ) {
            let rt = tokio::runtime::Runtime::new().unwrap();
            rt.block_on(async {
                let config = BackendConfig::cpu();
                let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config).await.unwrap();

                let tensor = Tensor::ones(&dims, &crate::Device::default()).unwrap();
                let mut test_tensor = tensor.clone();

                // Test that AllReduce works with any reasonable tensor shape
                let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
                prop_assert!(result.is_ok(), "AllReduce failed for dims {:?}", dims);

                pg.cleanup().await.unwrap();
                Ok(())
            });
        }

        #[test]
        fn prop_test_reduce_operations(
            op in prop_oneof![
                Just(ReduceOp::Sum),
                Just(ReduceOp::Max),
                Just(ReduceOp::Min),
                Just(ReduceOp::Product),
            ],
        ) {
            let rt = tokio::runtime::Runtime::new().unwrap();
            rt.block_on(async {
                let config = BackendConfig::cpu();
                let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config).await.unwrap();

                let tensor = Tensor::ones(&vec![5], &crate::Device::default()).unwrap();
                let mut test_tensor = tensor.clone();

                // Test that different reduce operations work
                let result = pg.all_reduce(&mut test_tensor, op).await;
                prop_assert!(result.is_ok(), "AllReduce failed for op {:?}", op);

                pg.cleanup().await.unwrap();
                Ok(())
            });
        }
    }
}

/// Integration tests that require actual CUDA/NCCL environment
#[cfg(test)]
mod integration_tests {
    use super::*;

    /// Test that requires actual NCCL initialization
    /// This test is only run when CUDA is available
    #[tokio::test]
    #[ignore = "requires CUDA environment"]
    async fn test_real_nccl_initialization() {
        // This test would only pass in an environment with CUDA and NCCL
        let config = BackendConfig::nccl();
        let world_info = WorldInfo::new(1, 0, Backend::Nccl);

        let result = ProcessGroup::new(Backend::Nccl, world_info);

        match result {
            Ok(pg) => {
                // Verify NCCL backend is actually being used
                #[cfg(feature = "nccl")]
                {
                    assert!(pg.has_nccl_backend());
                    assert!(pg.communication_primitive().unwrap().is_some());
                }

                // Test actual NCCL operations
                let tensor = Tensor::ones(&vec![1000], &crate::Device::Cuda(0)).unwrap();
                let mut test_tensor = tensor.clone();

                let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
                assert!(result.is_ok(), "NCCL AllReduce failed");

                pg.cleanup().await.unwrap();
            }
            Err(e) => {
                println!(
                    "NCCL initialization failed (expected in non-CUDA environment): {}",
                    e
                );
            }
        }
    }

    /// Test multi-GPU operations
    #[tokio::test]
    #[ignore = "requires multi-GPU CUDA environment"]
    async fn test_multi_gpu_operations() {
        // This would test actual multi-GPU NCCL operations
        // Only runs in environments with multiple GPUs

        for device_id in 0..2 {
            // Test with 2 GPUs
            let world_info = WorldInfo::new(2, device_id, Backend::Nccl);
            let pg = ProcessGroup::new(Backend::Nccl, world_info);

            if let Ok(pg) = pg {
                let device = crate::Device::Cuda(device_id as usize);
                let tensor = Tensor::ones(&vec![1024], &device).unwrap();
                let mut test_tensor = tensor.clone();

                // Test cross-GPU AllReduce
                let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
                assert!(
                    result.is_ok(),
                    "Multi-GPU AllReduce failed for device {}",
                    device_id
                );

                pg.cleanup().await.unwrap();
            }
        }
    }

    /// Benchmark NCCL operations vs custom implementations
    #[tokio::test]
    #[ignore = "performance benchmark"]
    async fn benchmark_nccl_vs_custom() {
        use std::time::Instant;

        let tensor_sizes = vec![1024, 10240, 102400, 1024000];

        for size in tensor_sizes {
            println!("Benchmarking tensor size: {}", size);

            // Test NCCL backend (if available)
            let world_info = WorldInfo::new(1, 0, Backend::Nccl);
            if let Ok(nccl_pg) = ProcessGroup::new(Backend::Nccl, world_info) {
                let tensor = Tensor::ones(&vec![size], &crate::Device::default()).unwrap();
                let mut nccl_tensor = tensor.clone();

                let start = Instant::now();
                let _ = nccl_pg.all_reduce(&mut nccl_tensor, ReduceOp::Sum).await;
                let nccl_duration = start.elapsed();

                println!("NCCL AllReduce time: {:?}", nccl_duration);
                nccl_pg.cleanup().await.unwrap();
            }

            // Test CPU backend for comparison
            let world_info = WorldInfo::new(1, 0, Backend::Cpu);
            let cpu_pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();
            let tensor = Tensor::ones(&vec![size], &crate::Device::default()).unwrap();
            let mut cpu_tensor = tensor.clone();

            let start = Instant::now();
            let _ = cpu_pg.all_reduce(&mut cpu_tensor, ReduceOp::Sum).await;
            let cpu_duration = start.elapsed();

            println!("CPU AllReduce time: {:?}", cpu_duration);
            cpu_pg.cleanup().await.unwrap();
        }
    }
}