//! Concurrency and Thread Safety Tests for arc-primitives
//!
//! This test suite validates thread-safe operation of cryptographic primitives.
//!
//! Test coverage:
//! - Parallel key generation (no race conditions)
//! - Concurrent encrypt/decrypt operations
//! - Thread-local RNG safety
//! - Lock-free operation verification
//! - Rayon parallel iterator compatibility
//!
//! These tests ensure the library is safe for enterprise multi-threaded deployments.

#![allow(clippy::expect_used)]

use latticearc::primitives::kem::ml_kem::{MlKem, MlKemSecurityLevel};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;

// ============================================================================
// Parallel Key Generation Tests
// ============================================================================

#[test]
fn test_parallel_keygen_produces_unique_keys_are_unique() {
    const NUM_THREADS: usize = 8;
    const KEYS_PER_THREAD: usize = 10;

    let keys = Arc::new(std::sync::Mutex::new(Vec::new()));

    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            let keys = Arc::clone(&keys);
            thread::spawn(move || {
                let mut local_keys = Vec::new();

                for _ in 0..KEYS_PER_THREAD {
                    let (pk, _sk) = MlKem::generate_keypair(MlKemSecurityLevel::MlKem512)
                        .expect("keypair generation should succeed");
                    local_keys.push(pk.to_bytes());
                }

                let mut keys_guard = keys.lock().expect("mutex should not be poisoned");
                keys_guard.extend(local_keys);
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    let keys_guard = keys.lock().expect("mutex should not be poisoned");
    let total_keys = keys_guard.len();
    assert_eq!(total_keys, NUM_THREADS * KEYS_PER_THREAD, "Should have generated all keys");

    // Verify all keys are unique
    let mut unique_keys: std::collections::HashSet<Vec<u8>> = std::collections::HashSet::new();
    for key in keys_guard.iter() {
        unique_keys.insert(key.clone());
    }
    assert_eq!(
        unique_keys.len(),
        total_keys,
        "All generated keys should be unique (no RNG collision)"
    );
}

#[test]
fn test_concurrent_keygen_different_security_levels_succeeds() {
    let success_count = Arc::new(AtomicUsize::new(0));
    let levels =
        [MlKemSecurityLevel::MlKem512, MlKemSecurityLevel::MlKem768, MlKemSecurityLevel::MlKem1024];

    let handles: Vec<_> = levels
        .iter()
        .map(|&level| {
            let success_count = Arc::clone(&success_count);
            thread::spawn(move || {
                for _ in 0..5 {
                    let result = MlKem::generate_keypair(level);
                    if result.is_ok() {
                        success_count.fetch_add(1, Ordering::SeqCst);
                    }
                }
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    assert_eq!(
        success_count.load(Ordering::SeqCst),
        15,
        "All 15 key generations (5 per level x 3 levels) should succeed"
    );
}

// ============================================================================
// Concurrent Encrypt/Decrypt Tests
// ============================================================================

#[test]
fn test_concurrent_encapsulation_same_key_succeeds() {
    let (pk, sk) = MlKem::generate_keypair(MlKemSecurityLevel::MlKem768)
        .expect("keypair generation should succeed");

    const NUM_THREADS: usize = 10;
    let pk = Arc::new(pk);
    let sk = Arc::new(sk);
    let success_count = Arc::new(AtomicUsize::new(0));

    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            let pk = Arc::clone(&pk);
            let sk = Arc::clone(&sk);
            let success_count = Arc::clone(&success_count);
            thread::spawn(move || {
                // Encapsulate with shared public key
                let (ss_enc, ct) = MlKem::encapsulate(&pk).expect("encapsulation should succeed");

                // Decapsulate with shared secret key (read-only, thread-safe)
                let ss_dec = MlKem::decapsulate(&sk, &ct).expect("decapsulation should succeed");

                // Verify shared secrets match
                if ss_enc.as_bytes() == ss_dec.as_bytes() {
                    success_count.fetch_add(1, Ordering::SeqCst);
                }
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    assert_eq!(
        success_count.load(Ordering::SeqCst),
        NUM_THREADS,
        "All concurrent encrypt/decrypt should produce matching shared secrets"
    );
}

#[test]
fn test_concurrent_different_ciphertexts_are_unique() {
    let (pk, _sk) = MlKem::generate_keypair(MlKemSecurityLevel::MlKem512)
        .expect("keypair generation should succeed");

    const NUM_THREADS: usize = 8;
    let pk = Arc::new(pk);
    let ciphertexts = Arc::new(std::sync::Mutex::new(Vec::new()));

    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            let pk = Arc::clone(&pk);
            let ciphertexts = Arc::clone(&ciphertexts);
            thread::spawn(move || {
                let (_ss, ct) = MlKem::encapsulate(&pk).expect("encapsulation should succeed");

                let mut cts = ciphertexts.lock().expect("mutex should not be poisoned");
                cts.push(ct.into_bytes());
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    let cts = ciphertexts.lock().expect("mutex should not be poisoned");
    let mut unique_cts: std::collections::HashSet<Vec<u8>> = std::collections::HashSet::new();
    for ct in cts.iter() {
        unique_cts.insert(ct.clone());
    }

    assert_eq!(
        unique_cts.len(),
        NUM_THREADS,
        "All ciphertexts should be unique (different randomness per thread)"
    );
}

// ============================================================================
// Thread-Local RNG Safety Tests
// ============================================================================

#[test]
fn test_thread_local_rng_isolation_succeeds() {
    // Verify that each thread's RNG produces independent randomness
    const NUM_THREADS: usize = 4;
    let shared_secrets = Arc::new(std::sync::Mutex::new(Vec::new()));

    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            let shared_secrets = Arc::clone(&shared_secrets);
            thread::spawn(move || {
                let (pk, _sk) = MlKem::generate_keypair(MlKemSecurityLevel::MlKem512)
                    .expect("keypair generation should succeed");

                let (ss, _ct) = MlKem::encapsulate(&pk).expect("encapsulation should succeed");

                let mut secrets = shared_secrets.lock().expect("mutex should not be poisoned");
                secrets.push(ss.as_bytes().to_vec());
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    let secrets = shared_secrets.lock().expect("mutex should not be poisoned");
    let mut unique_secrets: std::collections::HashSet<Vec<u8>> = std::collections::HashSet::new();
    for secret in secrets.iter() {
        unique_secrets.insert(secret.clone());
    }

    assert_eq!(
        unique_secrets.len(),
        NUM_THREADS,
        "Each thread should produce unique shared secret (RNG isolation)"
    );
}

// ============================================================================
// Stress Tests for Thread Safety
// ============================================================================

#[test]
fn test_high_concurrency_stress_succeeds() {
    const NUM_THREADS: usize = 32;
    const OPERATIONS_PER_THREAD: usize = 5;

    let operation_count = Arc::new(AtomicUsize::new(0));
    let error_count = Arc::new(AtomicUsize::new(0));

    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            let operation_count = Arc::clone(&operation_count);
            let error_count = Arc::clone(&error_count);
            thread::spawn(move || {
                for _ in 0..OPERATIONS_PER_THREAD {
                    // Full KEM cycle
                    match MlKem::generate_keypair(MlKemSecurityLevel::MlKem768) {
                        Ok((pk, sk)) => match MlKem::encapsulate(&pk) {
                            Ok((ss_enc, ct)) => match MlKem::decapsulate(&sk, &ct) {
                                Ok(ss_dec) => {
                                    if ss_enc.as_bytes() == ss_dec.as_bytes() {
                                        operation_count.fetch_add(1, Ordering::SeqCst);
                                    } else {
                                        error_count.fetch_add(1, Ordering::SeqCst);
                                    }
                                }
                                Err(_) => {
                                    error_count.fetch_add(1, Ordering::SeqCst);
                                }
                            },
                            Err(_) => {
                                error_count.fetch_add(1, Ordering::SeqCst);
                            }
                        },
                        Err(_) => {
                            error_count.fetch_add(1, Ordering::SeqCst);
                        }
                    }
                }
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    let successful = operation_count.load(Ordering::SeqCst);
    let errors = error_count.load(Ordering::SeqCst);
    let expected = NUM_THREADS * OPERATIONS_PER_THREAD;

    assert_eq!(errors, 0, "No errors should occur under high concurrency");
    assert_eq!(successful, expected, "All {} operations should complete successfully", expected);
}

// ============================================================================
// Thread Safety with Shared Read Access
// ============================================================================

#[test]
fn test_shared_public_key_concurrent_read_succeeds() {
    let (pk, _sk) = MlKem::generate_keypair(MlKemSecurityLevel::MlKem1024)
        .expect("keypair generation should succeed");

    // Public key should be safely readable from multiple threads
    let pk = Arc::new(pk);
    let pk_bytes_original = pk.to_bytes();
    let consistent_count = Arc::new(AtomicUsize::new(0));

    const NUM_THREADS: usize = 16;
    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            let pk = Arc::clone(&pk);
            let pk_bytes_original = pk_bytes_original.clone();
            let consistent_count = Arc::clone(&consistent_count);
            thread::spawn(move || {
                // Read public key bytes multiple times
                for _ in 0..10 {
                    let bytes = pk.to_bytes();
                    if bytes == pk_bytes_original {
                        consistent_count.fetch_add(1, Ordering::SeqCst);
                    }
                }
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    assert_eq!(
        consistent_count.load(Ordering::SeqCst),
        NUM_THREADS * 10,
        "All concurrent reads should return consistent public key bytes"
    );
}

// ============================================================================
// Memory Safety Under Concurrent Access
// ============================================================================

#[test]
fn test_no_data_races_during_operation_succeeds() {
    // This test verifies no undefined behavior from data races
    // by performing many concurrent operations and checking for panics

    const NUM_ITERATIONS: usize = 50;
    let panic_count = Arc::new(AtomicUsize::new(0));

    for _ in 0..NUM_ITERATIONS {
        let panic_count = Arc::clone(&panic_count);

        let handles: Vec<_> = (0..4)
            .map(|i| {
                let panic_count = Arc::clone(&panic_count);
                thread::spawn(move || {
                    let level = match i % 3 {
                        0 => MlKemSecurityLevel::MlKem512,
                        1 => MlKemSecurityLevel::MlKem768,
                        _ => MlKemSecurityLevel::MlKem1024,
                    };

                    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                        let (pk, sk) =
                            MlKem::generate_keypair(level).expect("keygen should succeed");
                        let (ss_enc, ct) = MlKem::encapsulate(&pk).expect("encap should succeed");
                        let ss_dec = MlKem::decapsulate(&sk, &ct).expect("decap should succeed");
                        assert_eq!(ss_enc.as_bytes(), ss_dec.as_bytes());
                    })) {
                        Ok(()) => {}
                        Err(_) => {
                            panic_count.fetch_add(1, Ordering::SeqCst);
                        }
                    }
                })
            })
            .collect();

        for handle in handles {
            let _ = handle.join();
        }
    }

    assert_eq!(
        panic_count.load(Ordering::SeqCst),
        0,
        "No panics should occur under concurrent operations"
    );
}

// ============================================================================
// Performance Under Concurrency (Sanity Check)
// ============================================================================

#[test]
fn test_concurrent_operations_complete_in_reasonable_time_succeeds() {
    use std::time::Instant;

    const NUM_THREADS: usize = 4;
    const OPS_PER_THREAD: usize = 5;
    const MAX_DURATION_SECS: u64 = 30; // Should complete well under this

    let start = Instant::now();

    let handles: Vec<_> = (0..NUM_THREADS)
        .map(|_| {
            thread::spawn(move || {
                for _ in 0..OPS_PER_THREAD {
                    let (pk, sk) =
                        MlKem::generate_keypair(MlKemSecurityLevel::MlKem768).expect("keygen");
                    let (_ss, ct) = MlKem::encapsulate(&pk).expect("encap");
                    let _ = MlKem::decapsulate(&sk, &ct).expect("decap");
                }
            })
        })
        .collect();

    for handle in handles {
        handle.join().expect("thread should not panic");
    }

    let duration = start.elapsed();
    assert!(
        duration.as_secs() < MAX_DURATION_SECS,
        "Concurrent operations should complete in reasonable time (took {:?})",
        duration
    );
}
