//! Distributed training benchmarks

use crate::{
    BenchmarkConfig, BenchmarkMeasurement, BenchmarkResults, BenchmarkStatistics,
    time_benchmark_async,
};
use anyhow::Result;
use std::collections::HashMap;
use tracing::info;

pub struct DistributedTrainingBenchmarks;

impl DistributedTrainingBenchmarks {
    pub fn new() -> Self {
        Self
    }

    pub async fn run_all(&self, config: &BenchmarkConfig) -> Result<BenchmarkResults> {
        info!("Starting distributed training benchmarks");

        let mut results = BenchmarkResults::new("DistributedTraining".to_string(), config.clone());
        let mut all_measurements: HashMap<String, Vec<BenchmarkMeasurement>> = HashMap::new();

        // Multi-GPU scaling benchmarks
        all_measurements.insert(
            "multi_gpu_scaling".to_string(),
            self.benchmark_multi_gpu_scaling(config).await?,
        );

        // Gradient synchronization
        all_measurements.insert(
            "gradient_sync".to_string(),
            self.benchmark_gradient_sync(config).await?,
        );

        // Calculate statistics
        for (name, measurements) in all_measurements {
            match BenchmarkStatistics::from_measurements(&measurements) {
                Ok(stats) => results.add_benchmark(name, stats),
                Err(e) => results.add_warning(format!("Failed to calculate stats: {e}")),
            }
        }

        results.mark_complete();
        Ok(results)
    }

    async fn benchmark_multi_gpu_scaling(
        &self,
        config: &BenchmarkConfig,
    ) -> Result<Vec<BenchmarkMeasurement>> {
        let mut measurements = Vec::with_capacity(config.measurement_iterations);

        for _ in 0..config.measurement_iterations {
            let measurement = time_benchmark_async("multi_gpu_scaling".to_string(), || async {
                self.multi_gpu_training_step().await?;
                Ok(())
            })
            .await?;
            measurements.push(measurement);
        }

        Ok(measurements)
    }

    async fn benchmark_gradient_sync(
        &self,
        config: &BenchmarkConfig,
    ) -> Result<Vec<BenchmarkMeasurement>> {
        let mut measurements = Vec::with_capacity(config.measurement_iterations);

        for _ in 0..config.measurement_iterations {
            let measurement = time_benchmark_async("gradient_sync".to_string(), || async {
                self.gradient_synchronization().await?;
                Ok(())
            })
            .await?;
            measurements.push(measurement);
        }

        Ok(measurements)
    }

    async fn multi_gpu_training_step(&self) -> Result<()> {
        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        Ok(())
    }

    async fn gradient_synchronization(&self) -> Result<()> {
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        Ok(())
    }
}

impl Default for DistributedTrainingBenchmarks {
    fn default() -> Self {
        Self::new()
    }
}