use criterion::{criterion_group, criterion_main, Criterion, Throughput};
use std::hint::black_box;
use std::time::Duration;

use x25519_dalek::{PublicKey, StaticSecret};
use b_wing::kem::KWing;

// Number of random inputs to pre-generate.
// This prevents system RNG overhead from polluting the benchmark.
const BATCH_SIZE: usize = 1000;

fn bench_k_wing(c: &mut Criterion) {
    let mut group = c.benchmark_group("KWing");
    group.measurement_time(Duration::from_secs(10));
    group.warm_up_time(Duration::from_secs(3));

    // ======================================================================
    // 1. Key Generation (Public Key Derivation)
    // ======================================================================
    let mut keygen_seeds = vec![[0u8; 128]; BATCH_SIZE];
    for seed in &mut keygen_seeds {
        getrandom::fill(seed).unwrap();
    }

    group.throughput(Throughput::Elements(1));
    group.bench_function("KeyGen", |b| {
        let mut i = 0;
        b.iter(|| {
            let seed = &keygen_seeds[i % BATCH_SIZE];
            i += 1;
            let binding = KWing::from_seed(black_box(seed)).unwrap();
            let pk = binding.get_pub_key();
            black_box(pk);
        })
    });

    // ======================================================================
    // 2. Encapsulation
    // ======================================================================
    let mut static_seed = [0u8; 128];
    getrandom::fill(&mut static_seed).unwrap();
    let binding = KWing::from_seed(&static_seed).unwrap();
    let pk = binding.get_pub_key();

    let mut encaps_seeds = vec![[0u8; 128]; BATCH_SIZE];
    for seed in &mut encaps_seeds {
        getrandom::fill(seed).unwrap();
    }

    group.bench_function("Encapsulate", |b| {
        let mut i = 0;
        b.iter(|| {
            let seed = &encaps_seeds[i % BATCH_SIZE];
            i += 1;
            let res = KWing::encapsulate(black_box(seed), black_box(&pk)).unwrap();
            black_box(res);
        })
    });

    // ======================================================================
    // 3. Decapsulation
    // ======================================================================
    let mut decaps_seed = [0u8; 128];
    getrandom::fill(&mut decaps_seed).unwrap();
    let binding = KWing::from_seed(&decaps_seed).unwrap();
    let pk_decaps = binding.get_pub_key();

    let mut eph_seed = [0u8; 128];
    getrandom::fill(&mut eph_seed).unwrap();
    let (ct, _) = KWing::encapsulate(&eph_seed, &pk_decaps).unwrap();

    group.bench_function("Decapsulate", |b| {
        b.iter(|| {
            let binding = KWing::from_seed(black_box(&decaps_seed)).unwrap();
            let res = binding.decapsulate(black_box(&ct)).unwrap();
            black_box(res);
        })
    });

    group.finish();
}

fn bench_pure_x25519(c: &mut Criterion) {
    let mut group = c.benchmark_group("Baseline_X25519");
    group.measurement_time(Duration::from_secs(5));
    group.warm_up_time(Duration::from_secs(1));

    // ======================================================================
    // Pure X25519 Round Trip
    // ======================================================================
    group.bench_function("Round_Trip_Total", |b| {
        let mut seed_a = [0u8; 32];
        let mut seed_b = [0u8; 32];
        getrandom::fill(&mut seed_a).unwrap();
        getrandom::fill(&mut seed_b).unwrap();

        b.iter(|| {
            // 1. Alice KeyGen
            let alice_secret = StaticSecret::from(black_box(seed_a));
            let alice_pub = PublicKey::from(&alice_secret);

            // 2. Bob "Encapsulate"
            let bob_secret = StaticSecret::from(black_box(seed_b));
            let bob_pub = PublicKey::from(&bob_secret);
            let ss_bob = bob_secret.diffie_hellman(&alice_pub);

            // 3. Alice "Decapsulate"
            let ss_alice = alice_secret.diffie_hellman(&bob_pub);

            black_box((ss_alice, ss_bob));
        })
    });

    group.finish();
}

criterion_group!(benches, bench_k_wing, bench_pure_x25519);
criterion_main!(benches);
