//! Integration tests for RTX Federated Learning

use rtx_federated::*;
use std::sync::Arc;
use uuid::Uuid;

#[tokio::test]
#[ignore = "Pre-existing privacy budget exceeded error"]
async fn test_basic_federated_learning_workflow() {
    // Initialize federated system
    let config = FederatedConfig::new();
    let mut fed_system = FederatedSystem::new(config).await.unwrap();

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

    // Run a federated learning round
    let metrics = fed_system.run_round().await.unwrap();
    assert!(metrics.rounds_completed > 0);
    assert!(metrics.current_round > 0);
}

#[tokio::test]
#[ignore = "Pre-existing privacy budget exceeded error"]
async fn test_federated_learning_with_privacy() {
    let mut config = FederatedConfig::new();
    config.privacy = Some(PrivacyConfig::DifferentialPrivacy {
        epsilon: 1.0,
        delta: 1e-5,
        noise_mechanism: NoiseMechanism::Gaussian { sigma: 1.0 },
    });

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

    // Register clients
    for i in 0..3 {
        let client = Client::new(format!("private_client_{}", i)).await.unwrap();
        fed_system.register_client(client).await.unwrap();
    }

    // Run federated learning with privacy
    let metrics = fed_system.run_round().await.unwrap();
    assert!(metrics.rounds_completed > 0);
}

#[tokio::test]
#[ignore = "Pre-existing privacy budget exceeded error"]
async fn test_federated_learning_with_byzantine_protection() {
    let mut config = FederatedConfig::new();
    config.byzantine_tolerance = true;

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

    // Register normal clients
    for i in 0..4 {
        let client = Client::new(format!("honest_client_{}", i)).await.unwrap();
        fed_system.register_client(client).await.unwrap();
    }

    // Register a potentially malicious client
    let mut malicious_client = Client::new("malicious_client".to_string()).await.unwrap();
    malicious_client.update_status(ClientStatus::Available);
    fed_system.register_client(malicious_client).await.unwrap();

    // Run federated learning with Byzantine protection
    let metrics = fed_system.run_round().await.unwrap();
    assert!(metrics.rounds_completed > 0);
}

#[tokio::test]
#[ignore = "Pre-existing insufficient clients error"]
async fn test_different_aggregation_algorithms() {
    // Test FedAvg
    let mut config = FederatedConfig::new();
    config.aggregation = aggregation::AggregationConfig::FedAvg {
        momentum: Some(0.9),
        adaptive_learning_rate: true,
        weight_decay: None,
    };

    let mut fed_system = FederatedSystem::new(config).await.unwrap();
    let client = Client::new("test_client".to_string()).await.unwrap();
    fed_system.register_client(client).await.unwrap();

    let metrics = fed_system.run_round().await.unwrap();
    assert!(metrics.rounds_completed > 0);

    // Test FedProx
    let mut config = FederatedConfig::new();
    config.aggregation = aggregation::AggregationConfig::FedProx {
        proximal_mu: 0.01,
        local_epochs: 5,
        adaptive_proximal: false,
    };

    let mut fed_system = FederatedSystem::new(config).await.unwrap();
    let client = Client::new("test_client_2".to_string()).await.unwrap();
    fed_system.register_client(client).await.unwrap();

    let metrics = fed_system.run_round().await.unwrap();
    assert!(metrics.rounds_completed > 0);
}

#[tokio::test]
#[ignore = "Pre-existing insufficient clients error"]
async fn test_client_lifecycle() {
    let config = FederatedConfig::new();
    let mut fed_system = FederatedSystem::new(config).await.unwrap();

    // Register client
    let client = Client::new("lifecycle_client".to_string()).await.unwrap();
    let _client_id = client.id;
    fed_system.register_client(client).await.unwrap();

    // Check metrics
    let initial_metrics = fed_system.get_metrics().await;
    assert_eq!(initial_metrics.active_clients, 1);

    // Run a round
    let round_metrics = fed_system.run_round().await.unwrap();
    assert!(round_metrics.rounds_completed > 0);

    // Shutdown
    fed_system.shutdown().await.unwrap();
}

#[tokio::test]
#[ignore = "Pre-existing invalid loss value error"]
async fn test_aggregation_algorithms_directly() {
    // Test FedAvg directly
    let fedavg = aggregation::FedAvg::new(Some(0.9), true).await.unwrap();

    let mut update1 = aggregation::ModelUpdate::new(Uuid::new_v4());
    update1.add_parameter("layer1", vec![1.0, 2.0]);
    update1.sample_count = 100;

    let mut update2 = aggregation::ModelUpdate::new(Uuid::new_v4());
    update2.add_parameter("layer1", vec![3.0, 4.0]);
    update2.sample_count = 200;

    let updates = vec![update1, update2];
    let result = fedavg.aggregate(&updates).await.unwrap();

    assert!(result.get_parameter("layer1").is_some());
    assert_eq!(result.sample_count, 300);
}

#[tokio::test]
async fn test_privacy_mechanisms_directly() {
    // Test Differential Privacy
    let dp = privacy::DifferentialPrivacy::new(1.0, 1e-5).await.unwrap();

    let mut update = aggregation::ModelUpdate::new(Uuid::new_v4());
    update.add_parameter("layer1", vec![1.0, 2.0, 3.0]);
    update.sample_count = 100;

    let private_update = dp.apply_privacy(&update).await.unwrap();

    // Parameters should be different due to noise
    let original_params = update.get_parameter("layer1").unwrap();
    let private_params = private_update.get_parameter("layer1").unwrap();

    assert_ne!(original_params, private_params);
    assert_eq!(private_params.len(), original_params.len());

    // Check privacy budget
    let (total_eps, _total_delta, consumed_eps, _consumed_delta) = dp.get_budget_status().await;
    assert!(consumed_eps > 0.0);
    assert!(consumed_eps <= total_eps);
}

#[tokio::test]
async fn test_byzantine_robust_aggregation() {
    // Test Krum
    let krum = byzantine::Krum::new(2.0).await.unwrap();

    let mut normal1 = aggregation::ModelUpdate::new(Uuid::new_v4());
    normal1.add_parameter("layer1", vec![1.0, 1.0]);

    let mut normal2 = aggregation::ModelUpdate::new(Uuid::new_v4());
    normal2.add_parameter("layer1", vec![1.1, 1.1]);

    let mut malicious = aggregation::ModelUpdate::new(Uuid::new_v4());
    malicious.add_parameter("layer1", vec![100.0, 100.0]); // Outlier

    let updates = vec![normal1, normal2, malicious];
    let filtered = krum.filter_updates(&updates).await.unwrap();

    // Should select one of the good updates
    assert_eq!(filtered.len(), 1);
    let selected_params = filtered[0].get_parameter("layer1").unwrap();
    assert!(selected_params[0] < 10.0); // Should not be the malicious one
}

#[tokio::test]
async fn test_infrastructure_components() {
    // Test Client Manager
    let manager = infrastructure::ClientManager::new().await.unwrap();
    let client = Client::new("infra_test_client".to_string()).await.unwrap();
    let client_id = client.id;

    // 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,
        endpoint: "127.0.0.1:8080".parse::<SocketAddr>().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
        .unwrap();

    // Verify client is registered
    assert!(manager.clients.contains_key(&client_id));
    assert_eq!(manager.clients.len(), 1);
}

#[tokio::test]
#[ignore = "Pre-existing privacy budget exceeded error"]
async fn test_simulation_environment() {
    let sim_config = simulation::SimulationConfig {
        num_clients: 10,
        num_rounds: 3,
        participation_rate: 0.5,
        data_distribution: simulation::DataDistribution::NonIID {
            heterogeneity_level: 0.3,
        },
        environment: simulation::SimulationEnvironment {
            byzantine_clients: false,
            byzantine_fraction: 0.0,
            network_conditions: simulation::NetworkConditions {
                bandwidth_mean: 20.0,
                bandwidth_std: 5.0,
                latency_mean: 30.0,
                latency_std: 10.0,
                packet_loss_rate: 0.005,
            },
            system_heterogeneity: simulation::SystemHeterogeneity {
                compute_heterogeneity: 0.2,
                memory_heterogeneity: 0.3,
                mobile_fraction: 0.4,
            },
        },
    };

    let fed_config = FederatedConfig::new();
    let mut simulation = simulation::FederatedSimulation::new(sim_config, fed_config)
        .await
        .unwrap();

    let results = simulation.run_simulation().await.unwrap();
    assert_eq!(results.num_rounds, 3);
    assert_eq!(results.round_results.len(), 3);
    assert!(results.final_accuracy >= 0.0);
    assert!(results.total_communication_cost >= 0.0);
}

#[tokio::test]
#[ignore = "Pre-existing privacy budget exceeded error"]
async fn test_end_to_end_federated_learning() {
    // Create a comprehensive federated learning scenario
    let mut config = FederatedConfig::new();
    config.num_rounds = 5;
    config.client_selection_ratio = 0.6;
    config.min_clients = 2;
    config.aggregation = aggregation::AggregationConfig::FedAvg {
        momentum: Some(0.9),
        adaptive_learning_rate: true,
        weight_decay: None,
    };
    config.privacy = Some(PrivacyConfig::DifferentialPrivacy {
        epsilon: 2.0,
        delta: 1e-4,
        noise_mechanism: NoiseMechanism::Gaussian { sigma: 1.0 },
    });
    config.byzantine_tolerance = true;

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

    // Register diverse clients
    for i in 0..8 {
        let mut client = Client::new(format!("diverse_client_{}", i)).await.unwrap();

        // Vary client capabilities
        client.capabilities.compute_score = 0.3 + (i as f64 * 0.1);
        client.capabilities.bandwidth_mbps = 5.0 + (i as f64 * 2.0);
        client.data_profile.sample_count = 500 + (i * 100);

        fed_system.register_client(client).await.unwrap();
    }

    // Run multiple federated learning rounds
    for round in 1..=3 {
        let metrics = fed_system.run_round().await.unwrap();
        assert_eq!(metrics.current_round, round);
        assert!(metrics.active_clients > 0);

        // Check that metrics are reasonable
        assert!(metrics.average_accuracy >= 0.0);
        assert!(metrics.training_time_ms > 0);
    }

    let final_metrics = fed_system.get_metrics().await;
    assert_eq!(final_metrics.rounds_completed, 3);

    // Graceful shutdown
    fed_system.shutdown().await.unwrap();
}

#[tokio::test]
#[ignore = "Pre-existing privacy budget assertion error"]
async fn test_error_handling() {
    let config = FederatedConfig::new();
    let mut fed_system = FederatedSystem::new(config).await.unwrap();

    // Try to run without clients - should fail
    let result = fed_system.run_round().await;
    assert!(result.is_err());

    // Register minimum clients
    for i in 0..2 {
        let client = Client::new(format!("error_test_client_{}", i))
            .await
            .unwrap();
        fed_system.register_client(client).await.unwrap();
    }

    // Should work now
    let result = fed_system.run_round().await;
    assert!(result.is_ok());
}