//! Example demonstrating NCCL backend usage for distributed training
//!
//! This example shows how to:
//! 1. Initialize a process group with NCCL backend
//! 2. Perform collective communication operations
//! 3. Handle different tensor operations
//! 4. Proper cleanup and error handling
//!
//! Usage:
//! ```bash
//! # Single GPU example
//! cargo run --example nccl_example --features nccl
//!
//! # Multi-GPU example (requires multiple GPUs)
//! CUDA_VISIBLE_DEVICES=0,1 mpirun -n 2 cargo run --example nccl_example --features nccl
//! ```

use anyhow::Result;
use rtx_distributed::{
    Backend, BackendConfig, Device, ProcessGroup, Tensor, TensorShape, WorldInfo,
    comm::{CommunicationPrimitive, ReduceOp},
};
use std::time::Instant;
use tracing::{info, warn};

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    info!("Starting NCCL backend example");

    // Check if NCCL feature is enabled
    #[cfg(not(feature = "nccl"))]
    {
        warn!("NCCL feature not enabled. Compile with --features nccl");
        return Ok(());
    }

    #[cfg(feature = "nccl")]
    {
        // Run the NCCL example
        run_nccl_example().await?;
    }

    Ok(())
}

#[cfg(feature = "nccl")]
async fn run_nccl_example() -> Result<()> {
    // Configuration
    let world_size = 1; // Single process example
    let rank = 0;

    info!("Initializing NCCL backend for rank {}/{}", rank, world_size);

    // Try NCCL backend first, fallback to CPU if CUDA not available
    let (pg, actual_backend) = match initialize_process_group(world_size, rank).await {
        Ok(pg) => {
            info!("Successfully initialized NCCL backend");
            (pg, "NCCL")
        }
        Err(e) => {
            warn!(
                "NCCL initialization failed: {}. Falling back to CPU backend",
                e
            );
            let config = BackendConfig::cpu();
            let pg =
                ProcessGroup::new_with_config(Backend::Cpu, world_size as i32, rank as i32, config)
                    .await?;
            (pg, "CPU")
        }
    };

    info!("Using {} backend", actual_backend);

    // Example 1: AllReduce operation
    info!("=== Example 1: AllReduce Operation ===");
    await_allreduce_example(&pg).await?;

    // Example 2: Broadcast operation
    info!("=== Example 2: Broadcast Operation ===");
    await_broadcast_example(&pg).await?;

    // Example 3: Performance comparison
    info!("=== Example 3: Performance Comparison ===");
    await_performance_comparison(&pg, actual_backend).await?;

    // Example 4: Different tensor shapes and data types
    info!("=== Example 4: Various Tensor Operations ===");
    await_tensor_operations_example(&pg).await?;

    // Example 5: Error handling
    info!("=== Example 5: Error Handling ===");
    await_error_handling_example(&pg).await?;

    // Cleanup
    info!("Cleaning up process group");
    pg.cleanup().await?;

    info!("NCCL backend example completed successfully");
    Ok(())
}

#[cfg(feature = "nccl")]
async fn initialize_process_group(world_size: usize, rank: usize) -> Result<ProcessGroup> {
    // Create NCCL backend configuration
    let mut config = BackendConfig::nccl();

    // Configure NCCL parameters for optimal performance
    config.set_parameter("nccl_socket_ifname".to_string(), "^lo".to_string());
    config.set_parameter("nccl_debug".to_string(), "INFO".to_string());
    config.set_timeout(std::time::Duration::from_secs(60));

    // Initialize process group
    Ok(
        ProcessGroup::new_with_config(Backend::Nccl, world_size as i32, rank as i32, config)
            .await?,
    )
}

#[cfg(feature = "nccl")]
async fn await_allreduce_example(pg: &ProcessGroup) -> Result<()> {
    // Create test tensor
    let shape = vec![1024];
    let mut tensor = Tensor::ones(&shape, &Device::default())?;

    info!(
        "Original tensor sum: {:.2}",
        tensor.data()?.iter().sum::<f32>()
    );

    // Perform AllReduce sum
    let start = Instant::now();
    pg.all_reduce(&mut tensor, ReduceOp::Sum).await?;
    let duration = start.elapsed();

    info!("AllReduce completed in {:?}", duration);
    info!(
        "Result tensor sum: {:.2}",
        tensor.data()?.iter().sum::<f32>()
    );

    // Test different reduce operations
    let mut max_tensor = Tensor::from_data(vec![1.0, 5.0, 3.0, 2.0], vec![4], &Device::default())?;

    pg.all_reduce(&mut max_tensor, ReduceOp::Max).await?;
    info!("Max reduce result: {:?}", max_tensor.data()?);

    Ok(())
}

#[cfg(feature = "nccl")]
async fn await_broadcast_example(pg: &ProcessGroup) -> Result<()> {
    // Create test tensor for broadcast
    let mut broadcast_tensor = Tensor::from_data(vec![42.0; 100], vec![100], &Device::default())?;

    info!("Broadcasting tensor with value 42.0");

    let start = Instant::now();
    pg.broadcast(&mut broadcast_tensor, 0).await?;
    let duration = start.elapsed();

    info!("Broadcast completed in {:?}", duration);
    info!("First few values: {:?}", &broadcast_tensor.data()?[0..5]);

    Ok(())
}

#[cfg(feature = "nccl")]
async fn await_performance_comparison(pg: &ProcessGroup, backend_name: &str) -> Result<()> {
    let sizes = vec![1024, 10240, 102400];

    info!("Performance comparison with {} backend:", backend_name);

    for size in sizes {
        let mut tensor = Tensor::ones(&vec![size], &Device::default())?;

        let start = Instant::now();
        pg.all_reduce(&mut tensor, ReduceOp::Sum).await?;
        let duration = start.elapsed();

        let bandwidth_gb = (size * 4) as f64 / (1024.0 * 1024.0 * 1024.0); // Approximate
        let bandwidth_gbps = bandwidth_gb / duration.as_secs_f64();

        info!(
            "Size: {} elements, Time: {:?}, Bandwidth: {:.2} GB/s",
            size, duration, bandwidth_gbps
        );
    }

    Ok(())
}

#[cfg(feature = "nccl")]
async fn await_tensor_operations_example(pg: &ProcessGroup) -> Result<()> {
    // Test different tensor shapes
    let test_shapes = vec![
        vec![1],
        vec![100],
        vec![10, 10],
        vec![5, 5, 5],
        vec![2, 2, 2, 2, 2],
    ];

    for shape_dims in test_shapes {
        info!("Testing shape: {:?}", shape_dims);

        let mut tensor = Tensor::ones(&shape_dims, &Device::default())?;
        let original_sum: f32 = tensor.data()?.iter().sum();

        // AllReduce
        pg.all_reduce(&mut tensor, ReduceOp::Sum).await?;
        let new_sum: f32 = tensor.data()?.iter().sum();

        info!(
            "  Original sum: {:.2}, After AllReduce: {:.2}",
            original_sum, new_sum
        );

        // Broadcast
        let mut broadcast_tensor =
            Tensor::zeros(TensorShape::new(shape_dims.clone())?, &Device::default())?;
        pg.broadcast(&mut broadcast_tensor, 0).await?;

        info!("  Broadcast completed for shape {:?}", shape_dims);
    }

    Ok(())
}

#[cfg(feature = "nccl")]
async fn await_error_handling_example(pg: &ProcessGroup) -> Result<()> {
    info!("Testing error conditions:");

    // Test broadcast with invalid root
    let mut tensor = Tensor::ones(&vec![10], &Device::default())?;
    match pg.broadcast(&mut tensor, 999).await {
        Ok(_) => warn!("Broadcast with invalid root should have failed"),
        Err(e) => info!("  ✓ Correctly caught invalid root error: {}", e),
    }

    // Test with very large tensor (might succeed or fail depending on memory)
    match Tensor::ones(&vec![100_000_000], &Device::default()) {
        Ok(mut large_tensor) => {
            info!("  Testing with large tensor (100M elements)");
            match pg.all_reduce(&mut large_tensor, ReduceOp::Sum).await {
                Ok(_) => info!("  ✓ Large tensor operation succeeded"),
                Err(e) => info!("  ✓ Large tensor operation failed as expected: {}", e),
            }
        }
        Err(e) => info!("  ✓ Large tensor creation failed as expected: {}", e),
    }

    Ok(())
}

#[cfg(not(feature = "nccl"))]
async fn run_nccl_example() -> Result<()> {
    warn!("This example requires the 'nccl' feature to be enabled");
    warn!("Please run with: cargo run --example nccl_example --features nccl");
    Ok(())
}