//! Standalone NCCL test to verify cudarc integration works
//!
//! This file tests NCCL functionality directly using cudarc
//! without depending on the rtx-distributed library.
//!
//! Run with: cargo run --bin nccl_test_standalone --features nccl

use std::sync::Arc;

#[cfg(feature = "nccl")]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    use cudarc::driver::CudaContext;
    use cudarc::nccl::{Id, Comm, ReduceOp as NcclReduceOp};

    println!("🚀 Starting NCCL integration test...");

    // Test 1: CUDA Context Creation
    println!("\n1️⃣  Testing CUDA context creation...");
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => {
            println!("✅ CUDA context created successfully on device {}", ctx.ordinal());
            Arc::new(ctx)
        }
        Err(e) => {
            println!("❌ CUDA context creation failed: {:?}", e);
            println!("💡 This might be expected if you don't have CUDA installed");
            return Ok(());
        }
    };

    // Test 2: NCCL ID Generation
    println!("\n2️⃣  Testing NCCL ID generation...");
    let id = match Id::new() {
        Ok(id) => {
            println!("✅ NCCL ID generated successfully");
            id
        }
        Err(e) => {
            println!("❌ NCCL ID generation failed: {:?}", e);
            println!("💡 This might be expected if you don't have NCCL installed");
            return Ok(());
        }
    };

    // Test 3: NCCL Communicator Creation
    println!("\n3️⃣  Testing NCCL communicator creation...");
    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => {
            println!("✅ NCCL communicator created successfully");
            println!("   Rank: {}, World size: {}", comm.rank(), comm.world_size());
            comm
        }
        Err(e) => {
            println!("❌ NCCL communicator creation failed: {:?}", e);
            return Err(e.into());
        }
    };

    // Test 4: Memory Allocation and AllReduce
    println!("\n4️⃣  Testing NCCL AllReduce operation...");
    let input_data = vec![1.0f32, 2.0, 3.0, 4.0];
    println!("   Input data: {:?}", input_data);

    // Allocate device memory
    let input_slice = stream.memcpy_stod(&input_data)?;
    let mut output_slice = stream.alloc_zeros::<f32>(input_data.len())?;

    // Perform AllReduce
    let reduce_op = NcclReduceOp::Sum;
    comm.all_reduce(&input_slice, &mut output_slice, &reduce_op)?;

    // Synchronize and get results
    stream.synchronize()?;
    let result_data = stream.memcpy_dtov(&output_slice)?;

    println!("   Output data: {:?}", result_data);

    // Verify results (for single rank, output should equal input)
    for (i, (&actual, &expected)) in result_data.iter().zip(input_data.iter()).enumerate() {
        if (actual - expected).abs() > 1e-6 {
            println!("❌ Result mismatch at index {}: expected {}, got {}", i, expected, actual);
            return Err("AllReduce result verification failed".into());
        }
    }
    println!("✅ AllReduce operation completed successfully");

    // Test 5: Broadcast Operation
    println!("\n5️⃣  Testing NCCL Broadcast operation...");
    let broadcast_data = vec![10.0f32, 20.0, 30.0, 40.0];
    println!("   Broadcast data: {:?}", broadcast_data);

    let broadcast_input = stream.memcpy_stod(&broadcast_data)?;
    let mut broadcast_output = stream.alloc_zeros::<f32>(broadcast_data.len())?;

    // For single rank, we're both sender and receiver
    comm.broadcast(Some(&broadcast_input), &mut broadcast_output, 0)?;

    stream.synchronize()?;
    let broadcast_result = stream.memcpy_dtov(&broadcast_output)?;

    println!("   Broadcast result: {:?}", broadcast_result);

    // Verify broadcast results
    for (i, (&actual, &expected)) in broadcast_result.iter().zip(broadcast_data.iter()).enumerate() {
        if (actual - expected).abs() > 1e-6 {
            println!("❌ Broadcast result mismatch at index {}: expected {}, got {}", i, expected, actual);
            return Err("Broadcast result verification failed".into());
        }
    }
    println!("✅ Broadcast operation completed successfully");

    println!("\n🎉 All NCCL tests passed! Integration is working correctly.");
    Ok(())
}

#[cfg(not(feature = "nccl"))]
fn main() {
    println!("⚠️  NCCL feature is not enabled.");
    println!("💡 Run with: cargo run --bin nccl_test_standalone --features nccl");
}