//! Load Generator for ML Training Service
//!
//! Configurable load generation for stress testing and performance analysis.

use anyhow::Result;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};

/// Load test configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoadTestConfig {
    /// Total number of operations to perform
    pub total_operations: u64,

    /// Number of concurrent workers
    pub concurrent_workers: usize,

    /// Operations per second (0 = unlimited)
    pub target_rps: u64,

    /// Test duration (None = run until total_operations complete)
    pub duration: Option<Duration>,

    /// Warmup duration before collecting metrics
    pub warmup_duration: Duration,

    /// Operation timeout
    pub operation_timeout: Duration,
}

impl Default for LoadTestConfig {
    fn default() -> Self {
        Self {
            total_operations: 1000,
            concurrent_workers: 10,
            target_rps: 0,
            duration: None,
            warmup_duration: Duration::from_secs(5),
            operation_timeout: Duration::from_secs(30),
        }
    }
}

/// Load test results
#[derive(Debug, Clone)]
pub struct LoadTestResults {
    pub total_operations: u64,
    pub successful_operations: u64,
    pub failed_operations: u64,
    pub total_duration: Duration,
    pub warmup_duration: Duration,
    pub actual_rps: f64,
    pub latencies: LatencyStats,
}

/// Latency statistics
#[derive(Debug, Clone)]
pub struct LatencyStats {
    pub min: Duration,
    pub max: Duration,
    pub mean: Duration,
    pub p50: Duration,
    pub p95: Duration,
    pub p99: Duration,
    pub p999: Duration,
}

/// Load generator
pub struct LoadGenerator<F, Fut>
where
    F: Fn(u64) -> Fut + Send + Sync + 'static,
    Fut: std::future::Future<Output = Result<()>> + Send,
{
    config: LoadTestConfig,
    operation: Arc<F>,
}

impl<F, Fut> LoadGenerator<F, Fut>
where
    F: Fn(u64) -> Fut + Send + Sync + 'static,
    Fut: std::future::Future<Output = Result<()>> + Send,
{
    /// Create a new load generator
    pub fn new(config: LoadTestConfig, operation: F) -> Self {
        Self {
            config,
            operation: Arc::new(operation),
        }
    }

    /// Run the load test
    pub async fn run(&self) -> Result<LoadTestResults> {
        println!("Starting load test with {} workers, target {} ops",
            self.config.concurrent_workers, self.config.total_operations);

        // Warmup phase
        println!("Warmup for {:?}...", self.config.warmup_duration);
        let warmup_start = Instant::now();
        self.warmup().await?;
        let warmup_duration = warmup_start.elapsed();

        // Main test phase
        println!("Starting main test phase...");
        let test_start = Instant::now();

        let completed = Arc::new(AtomicU64::new(0));
        let succeeded = Arc::new(AtomicU64::new(0));
        let failed = Arc::new(AtomicU64::new(0));
        let latencies = Arc::new(std::sync::Mutex::new(Vec::new()));

        let mut handles = Vec::new();

        // Calculate ops per worker
        let ops_per_worker = self.config.total_operations / self.config.concurrent_workers as u64;

        // Spawn worker tasks
        for worker_id in 0..self.config.concurrent_workers {
            let operation = self.operation.clone();
            let completed_count = completed.clone();
            let success_count = succeeded.clone();
            let fail_count = failed.clone();
            let latency_vec = latencies.clone();
            let timeout = self.config.operation_timeout;

            let handle = tokio::spawn(async move {
                for op_num in 0..ops_per_worker {
                    let global_op_num = worker_id as u64 * ops_per_worker + op_num;

                    let op_start = Instant::now();

                    match tokio::time::timeout(timeout, operation(global_op_num)).await {
                        Ok(Ok(_)) => {
                            success_count.fetch_add(1, Ordering::Relaxed);
                        }
                        Ok(Err(_)) | Err(_) => {
                            fail_count.fetch_add(1, Ordering::Relaxed);
                        }
                    }

                    let latency = op_start.elapsed();
                    latency_vec.lock().expect("INVARIANT: Lock should not be poisoned").push(latency);
                    completed_count.fetch_add(1, Ordering::Relaxed);
                }
            });

            handles.push(handle);
        }

        // Wait for all workers
        for handle in handles {
            handle.await?;
        }

        let total_duration = test_start.elapsed();

        // Calculate statistics
        let total_ops = completed.load(Ordering::Relaxed);
        let success_ops = succeeded.load(Ordering::Relaxed);
        let fail_ops = failed.load(Ordering::Relaxed);

        let actual_rps = total_ops as f64 / total_duration.as_secs_f64();

        let latency_stats = self.calculate_latency_stats(&latencies.lock().expect("INVARIANT: Lock should not be poisoned"));

        println!("✓ Load test complete");
        println!("  - Total ops: {}", total_ops);
        println!("  - Success: {}", success_ops);
        println!("  - Failed: {}", fail_ops);
        println!("  - Duration: {:?}", total_duration);
        println!("  - RPS: {:.2}", actual_rps);

        Ok(LoadTestResults {
            total_operations: total_ops,
            successful_operations: success_ops,
            failed_operations: fail_ops,
            total_duration,
            warmup_duration,
            actual_rps,
            latencies: latency_stats,
        })
    }

    /// Warmup phase
    async fn warmup(&self) -> Result<()> {
        let warmup_ops = 100; // Fixed warmup operation count
        let mut handles = Vec::new();

        for i in 0..warmup_ops {
            let operation = self.operation.clone();
            let handle = tokio::spawn(async move {
                let _ = operation(i).await;
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.await?;
        }

        tokio::time::sleep(self.config.warmup_duration).await;
        Ok(())
    }

    /// Calculate latency statistics
    fn calculate_latency_stats(&self, latencies: &[Duration]) -> LatencyStats {
        if latencies.is_empty() {
            return LatencyStats {
                min: Duration::ZERO,
                max: Duration::ZERO,
                mean: Duration::ZERO,
                p50: Duration::ZERO,
                p95: Duration::ZERO,
                p99: Duration::ZERO,
                p999: Duration::ZERO,
            };
        }

        let mut sorted = latencies.to_vec();
        sorted.sort();

        let len = sorted.len();
        let sum: Duration = sorted.iter().sum();

        LatencyStats {
            min: sorted[0],
            max: sorted[len - 1],
            mean: sum / len as u32,
            p50: sorted[len / 2],
            p95: sorted[(len * 95) / 100],
            p99: sorted[(len * 99) / 100],
            p999: sorted[(len * 999) / 1000],
        }
    }
}

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

    #[tokio::test]
    async fn test_load_generator_basic() -> Result<()> {
        let config = LoadTestConfig {
            total_operations: 100,
            concurrent_workers: 10,
            warmup_duration: Duration::from_millis(100),
            ..Default::default()
        };

        let operation = |_op_num: u64| async move {
            tokio::time::sleep(Duration::from_micros(100)).await;
            Ok(())
        };

        let generator = LoadGenerator::new(config, operation);
        let results = generator.run().await?;

        assert_eq!(results.total_operations, 100);
        assert_eq!(results.successful_operations, 100);
        assert_eq!(results.failed_operations, 0);

        Ok(())
    }

    #[tokio::test]
    async fn test_load_generator_with_failures() -> Result<()> {
        let config = LoadTestConfig {
            total_operations: 100,
            concurrent_workers: 10,
            warmup_duration: Duration::from_millis(100),
            ..Default::default()
        };

        let operation = |op_num: u64| async move {
            tokio::time::sleep(Duration::from_micros(100)).await;
            if op_num % 10 == 0 {
                Err(anyhow::anyhow!("Simulated failure"))
            } else {
                Ok(())
            }
        };

        let generator = LoadGenerator::new(config, operation);
        let results = generator.run().await?;

        assert_eq!(results.total_operations, 100);
        assert!(results.failed_operations >= 8, "Expected at least 8 failures");

        Ok(())
    }
}