//! Load generation for ML model stress testing

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;

/// Load generation profiles
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum LoadProfile {
    Constant,
    Ramp,
    Spike,
    Burst,
    Sine,
}

/// Traffic pattern configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrafficPattern {
    pub profile: LoadProfile,
    pub base_rps: u32,
    pub peak_rps: u32,
    pub pattern_duration_seconds: u64,
}

/// Load generator for stress testing
#[derive(Debug)]
pub struct LoadGenerator {
    target_rps: u32,
    concurrent_limit: Arc<Semaphore>,
    load_multiplier: f64,
}

impl LoadGenerator {
    pub fn new(target_rps: u32, concurrent_connections: u32) -> Result<Self> {
        Ok(Self {
            target_rps,
            concurrent_limit: Arc::new(Semaphore::new(concurrent_connections as usize)),
            load_multiplier: 1.0,
        })
    }

    pub fn set_load_multiplier(&mut self, multiplier: f64) {
        self.load_multiplier = multiplier;
    }

    pub fn get_current_rps(&self) -> u32 {
        (self.target_rps as f64 * self.load_multiplier) as u32
    }

    /// Generate load according to traffic pattern
    pub async fn generate_load<F>(&self, pattern: TrafficPattern, mut operation: F) -> Result<()>
    where
        F: FnMut() -> Result<()> + Send,
    {
        let start_time = Instant::now();
        let pattern_duration = Duration::from_secs(pattern.pattern_duration_seconds);

        while start_time.elapsed() < pattern_duration {
            let elapsed_ratio = start_time.elapsed().as_secs_f64() / pattern_duration.as_secs_f64();
            let current_rps = self.calculate_rps_for_pattern(&pattern, elapsed_ratio);

            let interval = Duration::from_secs_f64(1.0 / current_rps as f64);

            // Acquire semaphore permit for concurrency control
            let _permit = self.concurrent_limit.acquire().await?;

            // Execute operation
            operation()?;

            tokio::time::sleep(interval).await;
        }

        Ok(())
    }

    fn calculate_rps_for_pattern(&self, pattern: &TrafficPattern, elapsed_ratio: f64) -> u32 {
        let base = pattern.base_rps as f64;
        let peak = pattern.peak_rps as f64;

        let current_rps = match pattern.profile {
            LoadProfile::Constant => base,
            LoadProfile::Ramp => base + (peak - base) * elapsed_ratio,
            LoadProfile::Spike => {
                if elapsed_ratio > 0.8 && elapsed_ratio < 0.9 {
                    peak
                } else {
                    base
                }
            },
            LoadProfile::Burst => {
                if elapsed_ratio.rem_euclid(0.2) < 0.1 {
                    peak
                } else {
                    base
                }
            },
            LoadProfile::Sine => {
                base + (peak - base) * (std::f64::consts::PI * elapsed_ratio * 2.0).sin().abs()
            },
        };

        (current_rps * self.load_multiplier) as u32
    }
}
