#![cfg(unix)]

use std::hint::black_box;

use criterion::{criterion_group, criterion_main, Criterion, Throughput};
use rand::{rngs::StdRng, Rng, SeedableRng};
use std::time::Duration;

// Import the SecMem implementation
use sec_mem::{SecMem, zeroize::Zeroize};

fn benchmark_sec_mem(c: &mut Criterion) {
    let mut group = c.benchmark_group("SecMem");
    group.sample_size(500);
    group.measurement_time(Duration::from_secs(15));
    group.warm_up_time(Duration::from_secs(2));

    // Test different secret sizes
    let sizes = [16, 32, 64, 256, 1024, 4096];
    
    for &size in &sizes {
        group.throughput(Throughput::Bytes(size as u64));
        
        // Benchmark creation time
        group.bench_function(format!("create_{}b", size), |b| {
            let mut rng = StdRng::seed_from_u64(42);
            let mut secret = vec![0u8; size];
            rng.fill_bytes(&mut secret);
            
            b.iter(|| {
                let secret_copy = black_box(secret.clone());
                let _protected = SecMem::new(secret_copy);
                black_box(_protected);
            });
        });

        // Benchmark read access time
        group.bench_function(format!("read_access_{}b", size), |b| {
            let mut rng = StdRng::seed_from_u64(42);
            let mut secret = vec![0u8; size];
            rng.fill_bytes(&mut secret);
            let protected = SecMem::new(secret);
            
            b.iter(|| {
                protected.access(|s| {
                    let sum: u64 = s.iter().map(|&x| x as u64).sum();
                    black_box(sum);
                });
            });
        });

        // Benchmark write access time
        group.bench_function(format!("write_access_{}b", size), |b| {
            let mut rng = StdRng::seed_from_u64(42);
            let mut secret = vec![0u8; size];
            rng.fill_bytes(&mut secret);
            let mut protected = SecMem::new(secret);
            
            b.iter(|| {
                protected.access_mut(|s| {
                    for byte in s.iter_mut() {
                        *byte = black_box(*byte).wrapping_add(1);
                    }
                });
            });
        });

        // Benchmark constant-time comparison
        group.bench_function(format!("constant_time_eq_{}b", size), |b| {
            let mut rng = StdRng::seed_from_u64(42);
            let mut secret1 = vec![0u8; size];
            let mut secret2 = vec![0u8; size];
            rng.fill_bytes(&mut secret1);
            rng.fill_bytes(&mut secret2);
            
            let protected1 = SecMem::new(secret1.clone());
            let protected2 = SecMem::new(secret2);
            let protected1_same = SecMem::new(secret1);
            
            b.iter(|| {
                let result_diff = protected1.constant_time_eq(&protected2);
                let result_same = protected1.constant_time_eq(&protected1_same);
                black_box((result_diff, result_same));
            });
        });
    }
    
    group.finish();
}

fn benchmark_vs_unprotected(c: &mut Criterion) {
    let mut group = c.benchmark_group("SecMem_vs_Unprotected");
    group.sample_size(500);
    
    let size = 256; // Medium-sized secret
    let mut rng = StdRng::seed_from_u64(42);
    let mut secret_data = vec![0u8; size];
    rng.fill_bytes(&mut secret_data);
    
    // Benchmark unprotected access
    group.bench_function("unprotected_read", |b| {
        let secret = secret_data.clone();
        b.iter(|| {
            let sum: u64 = secret.iter().map(|&x| x as u64).sum();
            black_box(sum);
        });
    });
    
    // Benchmark protected access
    group.bench_function("protected_read", |b| {
        let protected = SecMem::new(secret_data.clone());
        b.iter(|| {
            protected.access(|s| {
                let sum: u64 = s.iter().map(|&x| x as u64).sum();
                black_box(sum);
            });
        });
    });
    
    // Benchmark unprotected write
    group.bench_function("unprotected_write", |b| {
        let mut secret = secret_data.clone();
        b.iter(|| {
            for byte in secret.iter_mut() {
                *byte = black_box(*byte).wrapping_add(1);
            }
            black_box(&mut secret);
        });
    });
    
    // Benchmark protected write
    group.bench_function("protected_write", |b| {
        let mut protected = SecMem::new(secret_data.clone());
        b.iter(|| {
            protected.access_mut(|s| {
                for byte in s.iter_mut() {
                    *byte = black_box(*byte).wrapping_add(1);
                }
            });
        });
    });
    
    group.finish();
}

fn benchmark_zeroize_performance(c: &mut Criterion) {
    let mut group = c.benchmark_group("Zeroize_Performance");
    
    let sizes = [64, 256, 1024];
    
    for &size in &sizes {
        group.bench_function(format!("zeroize_{}b", size), |b| {
            let mut rng = StdRng::seed_from_u64(42);
            let mut secret = vec![0u8; size];
            rng.fill_bytes(&mut secret);
            
            b.iter(|| {
                let mut protected = SecMem::new(secret.clone());
                protected.zeroize();
                black_box(protected);
            });
        });
    }
    
    group.finish();
}

fn benchmark_concurrent_access(c: &mut Criterion) {
    let mut group = c.benchmark_group("Concurrent_Access");
    
    let size = 128;
    let mut rng = StdRng::seed_from_u64(42);
    let mut secret = vec![0u8; size];
    rng.fill_bytes(&mut secret);
    let protected = std::sync::Arc::new(std::sync::Mutex::new(SecMem::new(secret)));
    
    group.bench_function("concurrent_reads", |b| {
        b.iter(|| {
            let protected_clone = protected.clone();
            let handles: Vec<_> = (0..4).map(|_| {
                let protected = protected_clone.clone();
                std::thread::spawn(move || {
                    protected.lock().unwrap().access(|s| {
                        let sum: u64 = s.iter().map(|&x| x as u64).sum();
                        black_box(sum);
                    })
                })
            }).collect();
            
            for handle in handles {
                handle.join().unwrap();
            }
        });
    });
    
    group.finish();
}

fn benchmark_memory_protection_overhead(c: &mut Criterion) {
    let mut group = c.benchmark_group("Memory_Protection_Overhead");
    
    let sizes = [64, 256, 1024];
    let iterations = 1000;
    
    for &size in &sizes {
        group.bench_function(format!("protection_cycles_{}b", size), |b| {
            let mut rng = StdRng::seed_from_u64(42);
            let mut secret = vec![0u8; size];
            rng.fill_bytes(&mut secret);
            let mut protected = SecMem::new(secret);
            
            b.iter(|| {
                for _ in 0..iterations {
                    protected.access_mut(|s| {
                        // Minimal operation to measure protection overhead
                        s[0] = s[0].wrapping_add(1);
                    });
                }
            });
        });
    }
    
    group.finish();
}

criterion_group!(
    benches,
    benchmark_sec_mem,
    benchmark_vs_unprotected,
    benchmark_zeroize_performance,
    benchmark_concurrent_access,
    benchmark_memory_protection_overhead
);

criterion_main!(benches);
