#![cfg(feature = "nccl")]
#![cfg(disabled)] // Disable until NCCL communicator APIs are fully implemented

#[allow(unused_imports)]
use cudarc::driver::CudaContext;
#[allow(unused_imports)]
use rtx_distributed::{Device, Result, comm::ReduceOp, error::DistributedError, nccl};
#[allow(unused_imports)]
use rtx_tensor::{DType, Shape, Tensor};
#[allow(unused_imports)]
use std::sync::Arc;
#[allow(unused_imports)]
use tokio::time::{Duration, timeout};

/// Integration tests for NCCL distributed operations
/// These tests define the expected behavior and drive implementation (TDD)
///
/// NOTE: These tests are currently disabled because they require the full
/// NCCL communicator API to be implemented.
#[cfg(test)]
mod nccl_integration_tests {
    use super::*;

    /// Test basic NCCL communicator creation and initialization
    #[tokio::test]
    async fn test_nccl_communicator_creation() -> Result<()> {
        // Skip if no CUDA devices available
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let world_size = 1;
        let rank = 0;

        // Create CUDA context
        let device = Arc::new(CudaContext::new(0).map_err(|e| {
            DistributedError::runtime(format!("Failed to create CUDA context: {:?}", e))
        })?);

        // Create NCCL ID for coordination
        let nccl_id = nccl::NcclCommunicator::get_unique_id()?;

        // Create communicator
        let config = nccl::NcclConfig::default();
        let comm = nccl::NcclCommunicator::new(world_size, rank, &nccl_id, device, config)?;

        // Verify communicator properties
        assert_eq!(comm.world_size(), world_size);
        assert_eq!(comm.rank(), rank);

        Ok(())
    }

    /// Test NCCL AllReduce operation with sum reduction
    #[tokio::test]
    async fn test_nccl_allreduce_sum() -> Result<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let device = Device::Cuda(0);
        let comm = create_test_communicator(device.clone())?;

        // Create test tensor
        let shape = vec![4, 4];
        let mut tensor = Tensor::ones(&shape, &Device::Cuda(0))?;

        // Perform AllReduce with Sum operation
        let result = timeout(
            Duration::from_secs(5),
            comm.allreduce(&mut tensor, ReduceOp::Sum),
        )
        .await;

        match result {
            Ok(Ok(())) => {
                // Verify tensor values (with world_size=1, values should remain 1.0)
                let data = tensor.to_vec1::<f32>()?;
                for &val in data.iter() {
                    assert!((val - 1.0).abs() < 1e-6, "Expected 1.0, got {}", val);
                }
            }
            Ok(Err(e)) => panic!("AllReduce failed: {:?}", e),
            Err(_) => panic!("AllReduce timed out"),
        }

        Ok(())
    }

    /// Test NCCL Broadcast operation
    #[tokio::test]
    async fn test_nccl_broadcast() -> Result<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let device = Device::Cuda(0);
        let comm = create_test_communicator(device.clone())?;

        // Create test tensor with specific values
        let shape = Shape::new(vec![3, 3])?;
        let mut tensor = Tensor::zeros(shape, DType::F32, &device)?;

        // Set root rank values
        if comm.rank() == 0 {
            let data: Vec<f32> = (0..9).map(|i| i as f32).collect();
            tensor.copy_from_cpu(&data)?;
        }

        // Perform Broadcast from rank 0
        let result = timeout(Duration::from_secs(5), comm.broadcast(&mut tensor, 0)).await;

        match result {
            Ok(Ok(())) => {
                // Verify all ranks have the broadcasted values
                let data = tensor.to_vec1::<f32>()?;
                for (i, &val) in data.iter().enumerate() {
                    assert!(
                        (val - i as f32).abs() < 1e-6,
                        "Expected {}, got {} at index {}",
                        i,
                        val,
                        i
                    );
                }
            }
            Ok(Err(e)) => panic!("Broadcast failed: {:?}", e),
            Err(_) => panic!("Broadcast timed out"),
        }

        Ok(())
    }

    /// Test NCCL AllGather operation
    #[tokio::test]
    async fn test_nccl_allgather() -> Result<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let device = Device::Cuda(0);
        let comm = create_test_communicator(device.clone())?;

        // Create input tensor with rank-specific values
        let input_shape = Shape::new(vec![2, 2])?;
        let mut input_tensor = Tensor::zeros(input_shape, DType::F32, &device)?;

        // Fill with rank-specific data
        let rank_data: Vec<f32> = (0..4)
            .map(|i| (comm.rank() as f32) * 10.0 + i as f32)
            .collect();
        input_tensor.copy_from_cpu(&rank_data)?;

        // Perform AllGather
        let result = timeout(Duration::from_secs(5), comm.allgather(&input_tensor)).await;

        match result {
            Ok(Ok(output_tensor)) => {
                // Verify output tensor shape (should be [world_size * 2, 2])
                let expected_shape = Shape::new(vec![comm.world_size() * 2, 2])?;
                assert_eq!(output_tensor.shape(), &expected_shape);

                // Verify gathered data
                let output_data = output_tensor.to_vec1::<f32>()?;
                assert_eq!(output_data.len(), comm.world_size() * 4);
            }
            Ok(Err(e)) => panic!("AllGather failed: {:?}", e),
            Err(_) => panic!("AllGather timed out"),
        }

        Ok(())
    }

    /// Test NCCL ReduceScatter operation
    #[tokio::test]
    async fn test_nccl_reduce_scatter() -> Result<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let device = Device::Cuda(0);
        let comm = create_test_communicator(device.clone())?;

        // Create input tensor
        let input_shape = Shape::new(vec![4, 2])?; // Will be scattered to [2, 2] per rank
        let mut input_tensor = Tensor::ones(input_shape, DType::F32, &device)?;

        // Perform ReduceScatter with Sum operation
        let result = timeout(
            Duration::from_secs(5),
            comm.reduce_scatter(&input_tensor, ReduceOp::Sum),
        )
        .await;

        match result {
            Ok(Ok(output_tensor)) => {
                // Verify output tensor shape
                let expected_shape = Shape::new(vec![2, 2])?;
                assert_eq!(output_tensor.shape(), &expected_shape);

                // Verify reduced values
                let output_data = output_tensor.to_vec1::<f32>()?;
                for &val in output_data.iter() {
                    assert!(
                        (val - comm.world_size() as f32).abs() < 1e-6,
                        "Expected {}, got {}",
                        comm.world_size(),
                        val
                    );
                }
            }
            Ok(Err(e)) => panic!("ReduceScatter failed: {:?}", e),
            Err(_) => panic!("ReduceScatter timed out"),
        }

        Ok(())
    }

    /// Test NCCL point-to-point send/recv operations
    #[tokio::test]
    async fn test_nccl_send_recv() -> Result<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let device = Device::Cuda(0);
        let comm = create_test_communicator(device.clone())?;

        if comm.world_size() < 2 {
            println!("Need at least 2 ranks for send/recv test, skipping");
            return Ok(());
        }

        let shape = Shape::new(vec![3, 3])?;

        if comm.rank() == 0 {
            // Sender: create tensor with specific values and send to rank 1
            let mut send_tensor = Tensor::zeros(shape, DType::F32, &device)?;
            let send_data: Vec<f32> = (0..9).map(|i| i as f32 + 100.0).collect();
            send_tensor.copy_from_cpu(&send_data)?;

            let result = timeout(Duration::from_secs(5), comm.send(&send_tensor, 1)).await;

            match result {
                Ok(Ok(())) => (),
                Ok(Err(e)) => panic!("Send failed: {:?}", e),
                Err(_) => panic!("Send timed out"),
            }
        } else if comm.rank() == 1 {
            // Receiver: receive tensor from rank 0
            let mut recv_tensor = Tensor::zeros(shape, DType::F32, &device)?;

            let result = timeout(Duration::from_secs(5), comm.recv(&mut recv_tensor, 0)).await;

            match result {
                Ok(Ok(())) => {
                    // Verify received data
                    let recv_data = recv_tensor.to_vec1::<f32>()?;
                    for (i, &val) in recv_data.iter().enumerate() {
                        assert!(
                            (val - (i as f32 + 100.0)).abs() < 1e-6,
                            "Expected {}, got {} at index {}",
                            i as f32 + 100.0,
                            val,
                            i
                        );
                    }
                }
                Ok(Err(e)) => panic!("Recv failed: {:?}", e),
                Err(_) => panic!("Recv timed out"),
            }
        }

        Ok(())
    }

    /// Test NCCL communicator synchronization
    #[tokio::test]
    async fn test_nccl_synchronization() -> Result<()> {
        if !cuda_available() {
            println!("CUDA not available, skipping NCCL tests");
            return Ok(());
        }

        let device = Device::Cuda(0);
        let comm = create_test_communicator(device.clone())?;

        // Test synchronization
        let result = timeout(Duration::from_secs(5), comm.synchronize()).await;

        match result {
            Ok(Ok(())) => (),
            Ok(Err(e)) => panic!("Synchronization failed: {:?}", e),
            Err(_) => panic!("Synchronization timed out"),
        }

        Ok(())
    }

    // Helper functions

    /// Check if CUDA is available for testing
    fn cuda_available() -> bool {
        // Try to create a CUDA context to check availability
        CudaContext::new(0).is_ok()
    }

    /// Create a test NCCL communicator for single-GPU testing
    fn create_test_communicator(_device: Device) -> Result<nccl::NcclCommunicator> {
        let world_size = 1;
        let rank = 0;

        // Create CUDA context
        let device = Arc::new(CudaContext::new(0).map_err(|e| {
            DistributedError::runtime(format!("Failed to create CUDA context: {:?}", e))
        })?);

        let nccl_id = nccl::NcclCommunicator::get_unique_id()?;
        let config = nccl::NcclConfig::default();

        nccl::NcclCommunicator::new(world_size, rank, &nccl_id, device, config)
    }
}