//! State Transition Storm Stress Tests
//!
//! Tests rapid state transitions under high concurrent load.
//! Validates database trigger performance, index efficiency, and deadlock prevention.

use anyhow::Result;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use uuid::Uuid;

use ml_training_service::orchestrator::{JobStatus, TrainingJob};
use ml::training_pipeline::ProductionTrainingConfig;

/// Helper to create a minimal training config
fn create_test_config() -> ProductionTrainingConfig {
    ProductionTrainingConfig::default()
}

/// Test 1: 1000 Concurrent Status Updates
///
/// Validates database can handle 1000 concurrent UPDATE operations
/// without deadlocks or excessive lock contention.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_1000_concurrent_status_updates() -> Result<()> {
    println!("\n=== Test 1: 1000 Concurrent Status Updates ===");

    let start = Instant::now();
    let jobs = Arc::new(RwLock::new(std::collections::HashMap::new()));

    // Create 100 jobs
    for i in 0..100 {
        let config = create_test_config();
        let job = TrainingJob::new(
            "DQN".to_string(),
            config,
            format!("Job {}", i),
            std::collections::HashMap::new(),
        );
        jobs.write().await.insert(job.id, job);
    }

    let success_count = Arc::new(AtomicU32::new(0));
    let conflict_count = Arc::new(AtomicU32::new(0));

    let mut handles = Vec::new();

    // Perform 1000 status updates (10 per job)
    for _ in 0..1000 {
        let jobs_clone = jobs.clone();
        let success = success_count.clone();
        let conflicts = conflict_count.clone();

        let handle = tokio::spawn(async move {
            // Pick random job
            let job_ids: Vec<Uuid> = {
                let jobs_read = jobs_clone.read().await;
                jobs_read.keys().copied().collect()
            };

            if let Some(&job_id) = job_ids.get(rand::random::<usize>() % job_ids.len()) {
                // Update status
                let mut jobs_write = jobs_clone.write().await;
                if let Some(job) = jobs_write.get_mut(&job_id) {
                    // Simulate status update
                    job.status = JobStatus::Running;
                    job.progress_percentage = rand::random::<f32>() * 100.0;
                    success.fetch_add(1, Ordering::Relaxed);
                } else {
                    conflicts.fetch_add(1, Ordering::Relaxed);
                }
            }
        });

        handles.push(handle);
    }

    // Wait for all updates
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let success = success_count.load(Ordering::Relaxed);
    let conflicts = conflict_count.load(Ordering::Relaxed);

    println!("✓ Test 1 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Successful updates: {}/1000", success);
    println!("  - Conflicts: {}", conflicts);
    println!("  - Throughput: {:.0} updates/sec", 1000.0 / elapsed.as_secs_f64());

    assert!(success >= 950, "Expected at least 950 successful updates, got {}", success);
    assert!(elapsed < Duration::from_secs(5), "Expected completion under 5s, got {:?}", elapsed);

    Ok(())
}

/// Test 2: Rapid State Transitions (Pending → Running → Completed)
///
/// Tests state machine integrity under rapid transitions.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_rapid_state_transitions() -> Result<()> {
    println!("\n=== Test 2: Rapid State Transitions ===");

    let start = Instant::now();
    let jobs = Arc::new(RwLock::new(std::collections::HashMap::new()));

    // Create 100 jobs in Pending state
    let mut job_ids = Vec::new();
    for i in 0..100 {
        let config = create_test_config();
        let job = TrainingJob::new(
            "DQN".to_string(),
            config,
            format!("Transition job {}", i),
            std::collections::HashMap::new(),
        );
        let job_id = job.id;
        jobs.write().await.insert(job_id, job);
        job_ids.push(job_id);
    }

    let completed_count = Arc::new(AtomicU32::new(0));

    let mut handles = Vec::new();

    // Transition each job: Pending → Running → Completed
    for job_id in job_ids {
        let jobs_clone = jobs.clone();
        let completed = completed_count.clone();

        let handle = tokio::spawn(async move {
            // Pending → Running
            {
                let mut jobs_write = jobs_clone.write().await;
                if let Some(job) = jobs_write.get_mut(&job_id) {
                    job.status = JobStatus::Running;
                    job.started_at = Some(chrono::Utc::now());
                }
            }

            tokio::time::sleep(Duration::from_micros(100)).await;

            // Running → Completed
            {
                let mut jobs_write = jobs_clone.write().await;
                if let Some(job) = jobs_write.get_mut(&job_id) {
                    job.status = JobStatus::Completed;
                    job.completed_at = Some(chrono::Utc::now());
                    job.progress_percentage = 100.0;
                    completed.fetch_add(1, Ordering::Relaxed);
                }
            }
        });

        handles.push(handle);
    }

    // Wait for all transitions
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let completed = completed_count.load(Ordering::Relaxed);

    // Verify final states
    let jobs_read = jobs.read().await;
    let completed_jobs = jobs_read.values()
        .filter(|j| j.status == JobStatus::Completed)
        .count();

    println!("✓ Test 2 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Completed transitions: {}/100", completed);
    println!("  - Final completed jobs: {}/100", completed_jobs);
    println!("  - Avg transition time: {:?}", elapsed / 100);

    assert_eq!(completed, 100, "Expected all 100 jobs to complete transitions");
    assert_eq!(completed_jobs, 100, "Expected all jobs in Completed state");

    Ok(())
}

/// Test 3: Database Trigger Performance (1000 Updates/Sec)
///
/// Simulates database triggers firing on status updates.
/// Validates trigger execution time stays within acceptable bounds.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_database_trigger_performance() -> Result<()> {
    println!("\n=== Test 3: Database Trigger Performance ===");

    let start = Instant::now();
    let trigger_execution_count = Arc::new(AtomicU32::new(0));
    let trigger_execution_time = Arc::new(std::sync::Mutex::new(Duration::ZERO));

    let mut handles = Vec::new();

    // Simulate 1000 updates with trigger execution
    for i in 0..1000 {
        let exec_count = trigger_execution_count.clone();
        let exec_time = trigger_execution_time.clone();

        let handle = tokio::spawn(async move {
            let trigger_start = Instant::now();

            // Simulate trigger logic (update audit log, send notification, etc.)
            let config = create_test_config();
            let mut job = TrainingJob::new(
                "DQN".to_string(),
                config,
                format!("Trigger test {}", i),
                std::collections::HashMap::new(),
            );

            job.status = JobStatus::Running;
            job.progress_percentage = (i % 100) as f32;

            // Simulate trigger overhead (1-5ms typical)
            tokio::time::sleep(Duration::from_micros(500)).await;

            let trigger_elapsed = trigger_start.elapsed();
            exec_count.fetch_add(1, Ordering::Relaxed);

            let mut total_time = exec_time.lock().expect("INVARIANT: Lock should not be poisoned");
            *total_time += trigger_elapsed;
        });

        handles.push(handle);
    }

    // Wait for all triggers
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let exec_count = trigger_execution_count.load(Ordering::Relaxed);
    let total_trigger_time = *trigger_execution_time.lock().expect("INVARIANT: Lock should not be poisoned");
    let avg_trigger_time = total_trigger_time / exec_count;
    let throughput = exec_count as f64 / elapsed.as_secs_f64();

    println!("✓ Test 3 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Trigger executions: {}/1000", exec_count);
    println!("  - Avg trigger time: {:?}", avg_trigger_time);
    println!("  - Throughput: {:.0} updates/sec", throughput);
    println!("  - Target: >1000 updates/sec");

    assert_eq!(exec_count, 1000, "Expected all 1000 triggers to execute");
    assert!(throughput >= 1000.0, "Throughput {:.0}/sec below 1000/sec target", throughput);
    assert!(avg_trigger_time < Duration::from_millis(5),
        "Avg trigger time {:?} exceeds 5ms", avg_trigger_time);

    Ok(())
}

/// Test 4: Index Performance Under Load
///
/// Validates that database indexes remain efficient under heavy load.
/// Tests common query patterns (filter by status, sort by created_at).
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_index_performance_under_load() -> Result<()> {
    println!("\n=== Test 4: Index Performance Under Load ===");

    let start = Instant::now();
    let jobs = Arc::new(RwLock::new(std::collections::HashMap::new()));

    // Create 1000 jobs with varying states
    for i in 0..1000 {
        let config = create_test_config();
        let mut job = TrainingJob::new(
            "DQN".to_string(),
            config,
            format!("Index test {}", i),
            std::collections::HashMap::new(),
        );

        // Distribute across states
        job.status = match i % 5 {
            0 => JobStatus::Pending,
            1 => JobStatus::Running,
            2 => JobStatus::Completed,
            3 => JobStatus::Failed,
            _ => JobStatus::Stopped,
        };

        jobs.write().await.insert(job.id, job);
    }

    let query_count = Arc::new(AtomicU32::new(0));
    let query_times = Arc::new(std::sync::Mutex::new(Vec::new()));

    let mut handles = Vec::new();

    // Perform 100 concurrent queries
    for _ in 0..100 {
        let jobs_clone = jobs.clone();
        let count = query_count.clone();
        let times = query_times.clone();

        let handle = tokio::spawn(async move {
            let query_start = Instant::now();

            // Simulate indexed query: filter by status
            let jobs_read = jobs_clone.read().await;
            let _running_jobs: Vec<_> = jobs_read.values()
                .filter(|j| j.status == JobStatus::Running)
                .collect();

            let query_elapsed = query_start.elapsed();
            count.fetch_add(1, Ordering::Relaxed);

            let mut times_vec = times.lock().expect("INVARIANT: Lock should not be poisoned");
            times_vec.push(query_elapsed);
        });

        handles.push(handle);
    }

    // Wait for all queries
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let queries = query_count.load(Ordering::Relaxed);

    let times_vec = query_times.lock().expect("INVARIANT: Lock should not be poisoned");
    let mut sorted_times = times_vec.clone();
    sorted_times.sort();

    let p50 = sorted_times[queries as usize / 2];
    let p95 = sorted_times[(queries as usize * 95) / 100];

    println!("✓ Test 4 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Queries executed: {}/100", queries);
    println!("  - Query P50: {:?}", p50);
    println!("  - Query P95: {:?}", p95);
    println!("  - Target P95: <10ms");

    assert_eq!(queries, 100, "Expected all 100 queries to complete");
    assert!(p95 < Duration::from_millis(10), "P95 {:?} exceeds 10ms target", p95);

    Ok(())
}

/// Test 5: Deadlock Prevention Validation
///
/// Tests that concurrent updates to related records don't cause deadlocks.
/// Validates lock ordering and transaction isolation.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_deadlock_prevention() -> Result<()> {
    println!("\n=== Test 5: Deadlock Prevention Validation ===");

    let start = Instant::now();
    let jobs = Arc::new(RwLock::new(std::collections::HashMap::new()));

    // Create 20 jobs
    let mut job_ids = Vec::new();
    for i in 0..20 {
        let config = create_test_config();
        let job = TrainingJob::new(
            "DQN".to_string(),
            config,
            format!("Deadlock test {}", i),
            std::collections::HashMap::new(),
        );
        let job_id = job.id;
        jobs.write().await.insert(job_id, job);
        job_ids.push(job_id);
    }

    let success_count = Arc::new(AtomicU32::new(0));
    let deadlock_count = Arc::new(AtomicU32::new(0));

    let mut handles = Vec::new();

    // Create 100 operations that update pairs of jobs
    for _ in 0..100 {
        let jobs_clone = jobs.clone();
        let job_ids_clone = job_ids.clone();
        let success = success_count.clone();
        let deadlocks = deadlock_count.clone();

        let handle = tokio::spawn(async move {
            // Pick two random jobs
            let idx1 = rand::random::<usize>() % job_ids_clone.len();
            let idx2 = rand::random::<usize>() % job_ids_clone.len();

            if idx1 == idx2 {
                return;
            }

            // Always acquire locks in consistent order (by UUID)
            let (first_id, second_id) = if job_ids_clone[idx1] < job_ids_clone[idx2] {
                (job_ids_clone[idx1], job_ids_clone[idx2])
            } else {
                (job_ids_clone[idx2], job_ids_clone[idx1])
            };

            // Acquire lock and update both jobs
            let operation = async {
                let mut jobs_write = jobs_clone.write().await;

                if let Some(job1) = jobs_write.get_mut(&first_id) {
                    job1.progress_percentage += 1.0;
                }

                if let Some(job2) = jobs_write.get_mut(&second_id) {
                    job2.progress_percentage += 1.0;
                }

                success.fetch_add(1, Ordering::Relaxed);
            };

            // Use timeout to detect potential deadlocks
            if tokio::time::timeout(Duration::from_secs(5), operation).await.is_err() {
                deadlocks.fetch_add(1, Ordering::Relaxed);
            }
        });

        handles.push(handle);
    }

    // Wait for all operations
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let success = success_count.load(Ordering::Relaxed);
    let deadlocks = deadlock_count.load(Ordering::Relaxed);

    println!("✓ Test 5 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Successful operations: {}", success);
    println!("  - Detected deadlocks: {}", deadlocks);
    println!("  - Deadlock-free: {}", deadlocks == 0);

    assert_eq!(deadlocks, 0, "Detected {} deadlocks - lock ordering failed", deadlocks);
    assert!(success >= 95, "Expected at least 95 successful operations, got {}", success);

    Ok(())
}