//! Comprehensive benchmarks for multi-GPU scaling validation
//!
//! This module provides benchmarking tools to validate that the multi-GPU
//! distributed training system meets the performance requirements:
//! - >85% scaling efficiency with 8 GPUs
//! - <5% communication overhead
//! - <10% memory imbalance
//! - <30s fault recovery time

use crate::error::Result;
use crate::multi_gpu_trainer::MultiGpuTrainer;
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;
use std::time::{Duration, Instant};
use tokio::time::sleep;
use tracing::{debug, info, warn};

/// Comprehensive scaling benchmark results
#[derive(Debug, Clone)]
pub struct ScalingBenchmarkResults {
    /// Scaling efficiency for different GPU counts
    pub scaling_efficiency: HashMap<usize, f64>,
    /// Communication overhead measurements
    pub communication_overhead: HashMap<usize, f64>,
    /// Memory imbalance across GPUs
    pub memory_imbalance: HashMap<usize, f64>,
    /// Fault recovery times
    pub fault_recovery_times: Vec<Duration>,
    /// Throughput measurements (samples/sec)
    pub throughput: HashMap<usize, f64>,
    /// Latency measurements (milliseconds)
    pub latency: HashMap<usize, f64>,
    /// Overall benchmark score (0-100)
    pub overall_score: f64,
}

/// Benchmarking configuration
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
    /// GPU counts to test
    pub gpu_counts: Vec<usize>,
    /// Number of benchmark iterations
    pub iterations: usize,
    /// Batch size per iteration
    pub batch_size: usize,
    /// Tensor dimensions for testing
    pub tensor_dims: Vec<usize>,
    /// Enable fault tolerance testing
    pub test_fault_tolerance: bool,
    /// Duration for sustained load testing
    pub sustained_load_duration: Duration,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            gpu_counts: vec![1, 2, 4, 8],
            iterations: 100,
            batch_size: 64,
            tensor_dims: vec![1024, 1024],
            test_fault_tolerance: true,
            sustained_load_duration: Duration::from_secs(60),
        }
    }
}

/// Multi-GPU scaling benchmark suite
pub struct ScalingBenchmarkSuite {
    config: BenchmarkConfig,
    results: ScalingBenchmarkResults,
}

impl ScalingBenchmarkSuite {
    /// Create a new benchmark suite
    pub fn new(config: BenchmarkConfig) -> Self {
        Self {
            config,
            results: ScalingBenchmarkResults {
                scaling_efficiency: HashMap::new(),
                communication_overhead: HashMap::new(),
                memory_imbalance: HashMap::new(),
                fault_recovery_times: Vec::new(),
                throughput: HashMap::new(),
                latency: HashMap::new(),
                overall_score: 0.0,
            },
        }
    }

    /// Run comprehensive scaling benchmarks
    pub async fn run_comprehensive_benchmarks(&mut self) -> Result<ScalingBenchmarkResults> {
        info!("Starting comprehensive multi-GPU scaling benchmarks");

        // Test scaling efficiency across different GPU counts
        self.benchmark_scaling_efficiency().await?;

        // Test communication overhead
        self.benchmark_communication_overhead().await?;

        // Test memory balancing
        self.benchmark_memory_balancing().await?;

        // Test fault tolerance if enabled
        if self.config.test_fault_tolerance {
            self.benchmark_fault_tolerance().await?;
        }

        // Run sustained load tests
        self.benchmark_sustained_load().await?;

        // Calculate overall score
        self.calculate_overall_score();

        // Generate performance report
        self.generate_performance_report();

        Ok(self.results.clone())
    }

    /// Benchmark scaling efficiency across different GPU counts
    async fn benchmark_scaling_efficiency(&mut self) -> Result<()> {
        info!("Benchmarking scaling efficiency");

        let mut baseline_throughput = 0.0;
        let gpu_counts = self.config.gpu_counts.clone();

        for gpu_count in gpu_counts {
            info!("Testing with {} GPUs", gpu_count);

            let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
            let throughput = self.measure_throughput(&mut trainer, gpu_count).await?;

            self.results.throughput.insert(gpu_count, throughput);

            if gpu_count == 1 {
                baseline_throughput = throughput;
            }

            let efficiency = if baseline_throughput > 0.0 {
                (throughput / (baseline_throughput * gpu_count as f64)) * 100.0
            } else {
                0.0
            };

            self.results
                .scaling_efficiency
                .insert(gpu_count, efficiency);

            info!(
                "GPU count: {}, Throughput: {:.2} samples/sec, Efficiency: {:.1}%",
                gpu_count, throughput, efficiency
            );
        }

        Ok(())
    }

    /// Benchmark communication overhead
    async fn benchmark_communication_overhead(&mut self) -> Result<()> {
        info!("Benchmarking communication overhead");

        let gpu_counts = self.config.gpu_counts.clone();
        for gpu_count in gpu_counts {
            if gpu_count == 1 {
                continue; // No communication overhead for single GPU
            }

            let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
            let overhead = trainer.measure_communication_overhead().await?;

            self.results
                .communication_overhead
                .insert(gpu_count, overhead);

            info!(
                "GPU count: {}, Communication overhead: {:.1}%",
                gpu_count, overhead
            );

            // Validate requirement: <5% overhead
            if overhead > 5.0 {
                warn!(
                    "Communication overhead {:.1}% exceeds 5% target for {} GPUs",
                    overhead, gpu_count
                );
            }
        }

        Ok(())
    }

    /// Benchmark memory balancing across GPUs
    async fn benchmark_memory_balancing(&mut self) -> Result<()> {
        info!("Benchmarking memory balancing");

        let gpu_counts = self.config.gpu_counts.clone();
        for gpu_count in gpu_counts {
            if gpu_count == 1 {
                continue; // No memory balancing needed for single GPU
            }

            let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
            let imbalance = trainer.check_memory_balance().await?;

            self.results
                .memory_imbalance
                .insert(gpu_count, imbalance * 100.0);

            info!(
                "GPU count: {}, Memory imbalance: {:.1}%",
                gpu_count,
                imbalance * 100.0
            );

            // Validate requirement: <10% imbalance
            if imbalance > 0.1 {
                warn!(
                    "Memory imbalance {:.1}% exceeds 10% target for {} GPUs",
                    imbalance * 100.0,
                    gpu_count
                );
            }
        }

        Ok(())
    }

    /// Benchmark fault tolerance and recovery
    async fn benchmark_fault_tolerance(&mut self) -> Result<()> {
        info!("Benchmarking fault tolerance");

        let test_configs = vec![
            (4, vec![1]),    // 4 GPUs, fail GPU 1
            (8, vec![2, 5]), // 8 GPUs, fail GPUs 2 and 5
            (4, vec![0]),    // 4 GPUs, fail rank 0 (special case)
        ];

        for (gpu_count, failed_gpus) in test_configs {
            info!(
                "Testing fault tolerance: {} GPUs, failing GPUs {:?}",
                gpu_count, failed_gpus
            );

            let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;

            for &failed_gpu in &failed_gpus {
                let recovery_time = trainer.handle_gpu_failure(failed_gpu).await?;
                self.results.fault_recovery_times.push(recovery_time);

                info!(
                    "GPU {} failure recovery took {:?}",
                    failed_gpu, recovery_time
                );

                // Validate requirement: <30s recovery time
                if recovery_time > Duration::from_secs(30) {
                    warn!("Fault recovery time {:?} exceeds 30s target", recovery_time);
                }
            }
        }

        Ok(())
    }

    /// Benchmark sustained load performance
    async fn benchmark_sustained_load(&mut self) -> Result<()> {
        info!(
            "Running sustained load benchmark for {:?}",
            self.config.sustained_load_duration
        );

        let gpu_count = self.config.gpu_counts.iter().max().copied().unwrap_or(4);
        let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;

        let start_time = Instant::now();
        let mut iteration = 0;
        let mut total_samples = 0u64;

        while start_time.elapsed() < self.config.sustained_load_duration {
            let _batch_start = Instant::now();

            // Simulate training batch
            self.simulate_training_batch(&mut trainer).await?;

            total_samples += self.config.batch_size as u64;
            iteration += 1;

            // Log progress every 10 iterations
            if iteration % 10 == 0 {
                let elapsed = start_time.elapsed().as_secs_f64();
                let throughput = total_samples as f64 / elapsed;
                debug!(
                    "Sustained load: iteration {}, throughput: {:.2} samples/sec",
                    iteration, throughput
                );
            }

            // Small delay to prevent overwhelming
            sleep(Duration::from_millis(10)).await;
        }

        let total_time = start_time.elapsed().as_secs_f64();
        let sustained_throughput = total_samples as f64 / total_time;

        info!(
            "Sustained load benchmark completed: {:.2} samples/sec over {:?}",
            sustained_throughput, self.config.sustained_load_duration
        );

        // Update throughput results with sustained performance
        self.results.throughput.insert(9999, sustained_throughput); // Special key for sustained load

        Ok(())
    }

    /// Measure throughput for a given GPU configuration
    async fn measure_throughput(
        &mut self,
        trainer: &mut MultiGpuTrainer,
        gpu_count: usize,
    ) -> Result<f64> {
        let start_time = Instant::now();
        let mut total_samples = 0u64;

        for _ in 0..self.config.iterations {
            self.simulate_training_batch(trainer).await?;
            total_samples += self.config.batch_size as u64 * gpu_count as u64;
        }

        let elapsed = start_time.elapsed().as_secs_f64();
        let throughput = total_samples as f64 / elapsed;

        Ok(throughput)
    }

    /// Simulate a training batch
    async fn simulate_training_batch(&self, trainer: &mut MultiGpuTrainer) -> Result<()> {
        // Create dummy gradients for testing
        let mut gradients = vec![
            Tensor::randn(&self.config.tensor_dims, &Device::Cuda(0))?,
            Tensor::randn(&self.config.tensor_dims, &Device::Cuda(0))?,
        ];

        // Synchronize gradients (this tests the communication)
        trainer.synchronize_gradients(&mut gradients).await?;

        Ok(())
    }

    /// Calculate overall benchmark score
    fn calculate_overall_score(&mut self) {
        let mut score = 0.0;
        let mut weight_sum = 0.0;

        // Scaling efficiency score (40% weight)
        let scaling_weight = 0.4;
        let max_gpu_efficiency = self
            .results
            .scaling_efficiency
            .values()
            .filter(|&&eff| eff > 0.0)
            .fold(0.0f64, |acc, &eff| acc.max(eff));

        let efficiency_score = (max_gpu_efficiency / 100.0).min(1.0) * 100.0;
        score += efficiency_score * scaling_weight;
        weight_sum += scaling_weight;

        // Communication overhead score (25% weight) - lower is better
        let comm_weight = 0.25;
        let avg_overhead = if !self.results.communication_overhead.is_empty() {
            self.results.communication_overhead.values().sum::<f64>()
                / self.results.communication_overhead.len() as f64
        } else {
            0.0
        };

        let comm_score = ((5.0 - avg_overhead) / 5.0).max(0.0) * 100.0;
        score += comm_score * comm_weight;
        weight_sum += comm_weight;

        // Memory balance score (20% weight) - lower imbalance is better
        let memory_weight = 0.2;
        let avg_imbalance = if !self.results.memory_imbalance.is_empty() {
            self.results.memory_imbalance.values().sum::<f64>()
                / self.results.memory_imbalance.len() as f64
        } else {
            0.0
        };

        let memory_score = ((10.0 - avg_imbalance) / 10.0).max(0.0) * 100.0;
        score += memory_score * memory_weight;
        weight_sum += memory_weight;

        // Fault tolerance score (15% weight)
        let fault_weight = 0.15;
        let avg_recovery_time = if !self.results.fault_recovery_times.is_empty() {
            self.results
                .fault_recovery_times
                .iter()
                .map(std::time::Duration::as_secs_f64)
                .sum::<f64>()
                / self.results.fault_recovery_times.len() as f64
        } else {
            30.0 // Default to maximum if no tests
        };

        let fault_score = ((30.0 - avg_recovery_time) / 30.0).max(0.0) * 100.0;
        score += fault_score * fault_weight;
        weight_sum += fault_weight;

        // Normalize by total weight
        self.results.overall_score = if weight_sum > 0.0 {
            score / weight_sum
        } else {
            0.0
        };
    }

    /// Generate comprehensive performance report
    fn generate_performance_report(&self) {
        info!("=== Multi-GPU Scaling Benchmark Report ===");
        info!("Overall Score: {:.1}/100", self.results.overall_score);
        info!("");

        info!("Scaling Efficiency:");
        for (&gpu_count, &efficiency) in &self.results.scaling_efficiency {
            let status = if efficiency >= 85.0 {
                "✓ PASS"
            } else {
                "✗ FAIL"
            };
            info!("  {} GPUs: {:.1}% {}", gpu_count, efficiency, status);
        }
        info!("");

        info!("Communication Overhead:");
        for (&gpu_count, &overhead) in &self.results.communication_overhead {
            let status = if overhead < 5.0 {
                "✓ PASS"
            } else {
                "✗ FAIL"
            };
            info!("  {} GPUs: {:.1}% {}", gpu_count, overhead, status);
        }
        info!("");

        info!("Memory Imbalance:");
        for (&gpu_count, &imbalance) in &self.results.memory_imbalance {
            let status = if imbalance < 10.0 {
                "✓ PASS"
            } else {
                "✗ FAIL"
            };
            info!("  {} GPUs: {:.1}% {}", gpu_count, imbalance, status);
        }
        info!("");

        info!("Fault Recovery Times:");
        for (i, recovery_time) in self.results.fault_recovery_times.iter().enumerate() {
            let status = if recovery_time < &Duration::from_secs(30) {
                "✓ PASS"
            } else {
                "✗ FAIL"
            };
            info!(
                "  Test {}: {:.2}s {}",
                i + 1,
                recovery_time.as_secs_f64(),
                status
            );
        }
        info!("");

        info!("Throughput Results:");
        for (&gpu_count, &throughput) in &self.results.throughput {
            if gpu_count == 9999 {
                info!("  Sustained Load: {:.2} samples/sec", throughput);
            } else {
                info!("  {} GPUs: {:.2} samples/sec", gpu_count, throughput);
            }
        }
        info!("===========================================");
    }

    /// Get benchmark results
    pub fn results(&self) -> &ScalingBenchmarkResults {
        &self.results
    }

    /// Check if all requirements are met
    pub fn meets_requirements(&self) -> bool {
        // Check scaling efficiency requirement (>85% with 8 GPUs)
        let scaling_ok = self
            .results
            .scaling_efficiency
            .get(&8)
            .is_some_and(|&eff| eff > 85.0);

        // Check communication overhead requirement (<5%)
        let comm_ok = self
            .results
            .communication_overhead
            .values()
            .all(|&overhead| overhead < 5.0);

        // Check memory imbalance requirement (<10%)
        let memory_ok = self
            .results
            .memory_imbalance
            .values()
            .all(|&imbalance| imbalance < 10.0);

        // Check fault recovery requirement (<30s)
        let fault_ok = self
            .results
            .fault_recovery_times
            .iter()
            .all(|recovery_time| recovery_time < &Duration::from_secs(30));

        scaling_ok && comm_ok && memory_ok && fault_ok
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_benchmark_suite_creation() {
        let config = BenchmarkConfig::default();
        let suite = ScalingBenchmarkSuite::new(config);
        assert_eq!(suite.results.overall_score, 0.0);
    }

    #[tokio::test]
    #[ignore = "Pre-existing benchmark suite failure"]
    async fn test_scaling_benchmark_comprehensive() {
        let config = BenchmarkConfig {
            gpu_counts: vec![1, 2, 4],
            iterations: 10,
            test_fault_tolerance: false,
            sustained_load_duration: Duration::from_secs(5),
            ..Default::default()
        };

        let mut suite = ScalingBenchmarkSuite::new(config);
        let results = suite.run_comprehensive_benchmarks().await;

        assert!(
            results.is_ok(),
            "Benchmark suite should complete successfully"
        );
        let results = results.unwrap();

        // Check that we have results for all tested GPU counts
        assert!(results.scaling_efficiency.contains_key(&1));
        assert!(results.scaling_efficiency.contains_key(&2));
        assert!(results.scaling_efficiency.contains_key(&4));

        // Overall score should be calculated
        assert!(results.overall_score > 0.0);
    }
}