//! RTX Federated Learning Demo
//!
//! This example demonstrates the key capabilities of RTX Federated Learning:
//! - Advanced aggregation algorithms (FedAvg, FedProx)
//! - Privacy-preserving mechanisms (Differential Privacy)
//! - Byzantine fault tolerance (Krum)
//! - Production infrastructure components
//! - Simulation environment

use rtx_federated::*;
use std::sync::Arc;
use tracing::{info, warn};

#[tokio::main]
async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    tracing_subscriber::fmt::init();

    info!("🚀 RTX Federated Learning Demo Starting");

    // Demo 1: Basic Federated Learning
    demo_basic_federated_learning().await?;

    // Demo 2: Privacy-Preserving Federated Learning
    demo_privacy_preserving_learning().await?;

    // Demo 3: Byzantine-Robust Federated Learning
    demo_byzantine_robust_learning().await?;

    // Demo 4: Advanced Aggregation Algorithms
    demo_advanced_aggregation().await?;

    // Demo 5: Simulation Environment
    demo_simulation_environment().await?;

    info!("✅ RTX Federated Learning Demo Complete");
    Ok(())
}

async fn demo_basic_federated_learning() -> std::result::Result<(), Box<dyn std::error::Error>> {
    info!("📚 Demo 1: Basic Federated Learning");

    // Create federated learning configuration
    let config = FederatedConfig::new().with_aggregation(aggregation::AggregationConfig::FedAvg {
        momentum: Some(0.9),
        adaptive_learning_rate: true,
        weight_decay: None,
    });

    // Initialize federated system
    let mut fed_system = FederatedSystem::new(config).await?;

    // Register clients with diverse characteristics
    for i in 0..10 {
        let mut client = Client::new(format!("basic_client_{}", i)).await?;

        // Vary client capabilities
        client.capabilities.compute_score = 0.3 + (i as f64 * 0.07);
        client.capabilities.bandwidth_mbps = 10.0 + (i as f64 * 2.0);
        client.data_profile.sample_count = 800 + (i * 50);

        fed_system.register_client(client).await?;
    }

    // Run federated learning rounds
    for round in 1..=5 {
        info!("🔄 Running federated learning round {}", round);
        let metrics = fed_system.run_round().await?;

        info!("📊 Round {} Results:", round);
        info!("   - Participating clients: {}", metrics.active_clients);
        info!("   - Average accuracy: {:.4}", metrics.average_accuracy);
        info!("   - Training time: {}ms", metrics.training_time_ms);
    }

    let final_metrics = fed_system.get_metrics().await;
    info!("🎯 Final Results:");
    info!(
        "   - Total rounds completed: {}",
        final_metrics.rounds_completed
    );
    info!(
        "   - Final average accuracy: {:.4}",
        final_metrics.average_accuracy
    );

    fed_system.shutdown().await?;
    info!("✅ Basic federated learning demo completed\n");
    Ok(())
}

async fn demo_privacy_preserving_learning() -> std::result::Result<(), Box<dyn std::error::Error>> {
    info!("🔒 Demo 2: Privacy-Preserving Federated Learning");

    // Configure with differential privacy
    let config = FederatedConfig::new().with_aggregation(aggregation::AggregationConfig::FedAvg {
        momentum: Some(0.9),
        adaptive_learning_rate: true,
        weight_decay: None,
    });

    let mut fed_system = FederatedSystem::new(config).await?;

    // Register privacy-conscious clients
    for i in 0..6 {
        let mut client = Client::new(format!("private_client_{}", i)).await?;
        client.data_profile.privacy_level = if i % 2 == 0 {
            PrivacyLevel::Confidential
        } else {
            PrivacyLevel::Restricted
        };

        fed_system.register_client(client).await?;
    }

    // Note: Differential privacy mechanism is configured in the system
    info!("🔒 Privacy mechanisms are active in the federated system");

    // Run privacy-preserving federated learning
    for round in 1..=3 {
        let metrics = fed_system.run_round().await?;
        info!(
            "🔒 Private Round {}: {} clients, accuracy: {:.4}",
            round, metrics.active_clients, metrics.average_accuracy
        );
    }

    fed_system.shutdown().await?;
    info!("✅ Privacy-preserving learning demo completed\n");
    Ok(())
}

async fn demo_byzantine_robust_learning() -> std::result::Result<(), Box<dyn std::error::Error>> {
    info!("🛡️  Demo 3: Byzantine-Robust Federated Learning");

    // Configure with Byzantine tolerance
    let mut config = FederatedConfig::new();
    config.byzantine_tolerance = true;

    let mut fed_system = FederatedSystem::new(config).await?;

    // Register honest clients
    for i in 0..8 {
        let mut client = Client::new(format!("honest_client_{}", i)).await?;
        client.status = ClientStatus::Available;
        fed_system.register_client(client).await?;
    }

    // Register potentially malicious clients
    for i in 0..2 {
        let mut malicious_client = Client::new(format!("suspicious_client_{}", i)).await?;
        malicious_client.status = ClientStatus::Available;
        fed_system.register_client(malicious_client).await?;
    }

    // Demonstrate Krum algorithm directly
    let krum = byzantine::Krum::new(0.2).await?;

    // Create test updates with one outlier
    let mut updates = Vec::new();
    for i in 0..5 {
        let mut update = aggregation::ModelUpdate::new(uuid::Uuid::new_v4());
        if i < 4 {
            update.add_parameter("layer1", vec![1.0 + i as f64 * 0.1, 2.0 + i as f64 * 0.1]);
        } else {
            // Malicious update
            update.add_parameter("layer1", vec![100.0, 100.0]);
        }
        update.sample_count = 100;
        updates.push(update);
    }

    let filtered_updates = krum.filter_updates(&updates).await?;
    info!(
        "🛡️  Krum filtered {} updates to {} honest updates",
        updates.len(),
        filtered_updates.len()
    );

    // Run Byzantine-robust federated learning
    for round in 1..=3 {
        let metrics = fed_system.run_round().await?;
        info!(
            "🛡️  Robust Round {}: {} clients, {} Byzantine attacks detected",
            round, metrics.active_clients, metrics.byzantine_attacks_detected
        );
    }

    fed_system.shutdown().await?;
    info!("✅ Byzantine-robust learning demo completed\n");
    Ok(())
}

async fn demo_advanced_aggregation() -> std::result::Result<(), Box<dyn std::error::Error>> {
    info!("🧠 Demo 4: Advanced Aggregation Algorithms");

    // Test different aggregation algorithms
    let algorithms = vec![
        (
            "FedAvg",
            aggregation::AggregationConfig::FedAvg {
                momentum: Some(0.9),
                adaptive_learning_rate: true,
                weight_decay: None,
            },
        ),
        (
            "FedProx",
            aggregation::AggregationConfig::FedProx {
                proximal_mu: 0.01,
                local_epochs: 5,
                adaptive_proximal: false,
            },
        ),
        (
            "SCAFFOLD",
            aggregation::AggregationConfig::Scaffold {
                learning_rate: 0.01,
                local_steps: 100,
                variance_reduction: true,
            },
        ),
        (
            "FedNova",
            aggregation::AggregationConfig::FedNova {
                tau_effective: 10.0,
                momentum_factor: 0.9,
                normalize_weights: true,
            },
        ),
    ];

    for (name, algorithm_config) in algorithms {
        info!("🔬 Testing {} algorithm", name);

        let config = FederatedConfig::new().with_aggregation(algorithm_config);

        let mut fed_system = FederatedSystem::new(config).await?;

        // Register clients
        for i in 0..5 {
            let client = Client::new(format!("{}_client_{}", name.to_lowercase(), i)).await?;
            fed_system.register_client(client).await?;
        }

        // Run one round
        let metrics = fed_system.run_round().await?;
        info!(
            "   - {} result: {} clients, accuracy: {:.4}, time: {}ms",
            name, metrics.active_clients, metrics.average_accuracy, metrics.training_time_ms
        );

        fed_system.shutdown().await?;
    }

    info!("✅ Advanced aggregation algorithms demo completed\n");
    Ok(())
}

async fn demo_simulation_environment() -> std::result::Result<(), Box<dyn std::error::Error>> {
    info!("🎮 Demo 5: Simulation Environment");

    // Configure simulation
    let sim_config = simulation::SimulationConfig {
        num_clients: 50,
        num_rounds: 10,
        participation_rate: 0.3,
        data_distribution: simulation::DataDistribution::NonIID {
            heterogeneity_level: 0.4,
        },
        environment: simulation::SimulationEnvironment {
            byzantine_clients: true,
            byzantine_fraction: 0.1,
            network_conditions: simulation::NetworkConditions {
                bandwidth_mean: 15.0,
                bandwidth_std: 8.0,
                latency_mean: 45.0,
                latency_std: 15.0,
                packet_loss_rate: 0.02,
            },
            system_heterogeneity: simulation::SystemHeterogeneity {
                compute_heterogeneity: 0.4,
                memory_heterogeneity: 0.3,
                mobile_fraction: 0.6,
            },
        },
    };

    // Configure federated learning
    let mut fed_config = FederatedConfig::new();
    fed_config.client_selection_ratio = 0.3;
    fed_config.byzantine_tolerance = true;

    // Run simulation
    info!(
        "🎮 Starting simulation with {} clients for {} rounds",
        sim_config.num_clients, sim_config.num_rounds
    );

    let mut simulation = simulation::FederatedSimulation::new(sim_config, fed_config).await?;
    let results = simulation.run_simulation().await?;

    // Display results
    info!("🎯 Simulation Results:");
    info!("   - Total rounds: {}", results.num_rounds);
    info!("   - Final accuracy: {:.4}", results.final_accuracy);
    info!(
        "   - Total communication cost: {:.2}",
        results.total_communication_cost
    );
    info!("   - Convergence round: {:?}", results.convergence_round);

    // Show round-by-round progress
    info!("📈 Round-by-round progress:");
    for (i, round_result) in results.round_results.iter().enumerate() {
        if i % 2 == 0 || i == results.round_results.len() - 1 {
            info!(
                "   Round {}: {} participants, accuracy: {:.4}, cost: {:.2}",
                round_result.round_number,
                round_result.participating_clients,
                round_result.average_accuracy,
                round_result.communication_cost
            );
        }
    }

    info!("✅ Simulation environment demo completed\n");
    Ok(())
}

// Utility function to demonstrate infrastructure components
async fn _demo_infrastructure_components() -> std::result::Result<(), Box<dyn std::error::Error>> {
    info!("🏗️  Infrastructure Components Demo");

    // Client Manager
    let manager = infrastructure::ClientManager::new().await?;
    for i in 0..5 {
        let client = Client::new(format!("infra_client_{}", i)).await?;

        // Create ClientConnection from Client
        use infrastructure::client_manager::protocol::{ClientConnection, ConnectionQuality};
        use std::net::SocketAddr;
        use std::sync::atomic::AtomicU64;

        let connection = ClientConnection {
            client_id: client.id,
            endpoint: client
                .endpoint
                .as_ref()
                .and_then(|e| e.parse::<SocketAddr>().ok())
                .unwrap_or_else(|| "127.0.0.1:8080".parse().unwrap()),
            connected_at: chrono::Utc::now(),
            last_heartbeat: chrono::Utc::now(),
            bytes_sent: Arc::new(AtomicU64::new(0)),
            bytes_received: Arc::new(AtomicU64::new(0)),
            round_trip_time_ms: Arc::new(AtomicU64::new(0)),
            connection_quality: ConnectionQuality::default(),
            auth_token: None,
            stream: None,
        };

        manager.register_client(client.id, connection).await?;
    }

    let total_clients = manager.connected_clients();
    info!(
        "📊 Client Manager Stats: {} total clients connected",
        total_clients
    );

    info!("✅ Infrastructure components demo completed");
    Ok(())
}