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

use ed25519_dalek::{SigningKey, Signer, Verifier};
use b_wing::signing::SWing;

// Keep the batch size reasonable since SLH-DSA KeyGen is heavy
const BATCH_SIZE: usize = 100;
const TEST_MESSAGE: &[u8] = b"Paranoia is just a higher level of awareness.";
const TEST_CONTEXT: &[u8] = b"b-wing-benchmark-context";

fn bench_s_wing(c: &mut Criterion) {
    let mut group = c.benchmark_group("SWing");
    // Give Criterion a bit more time to gather samples because SLH-DSA is heavy
    group.measurement_time(Duration::from_secs(60));
    group.warm_up_time(Duration::from_secs(5));

    // ======================================================================
    // 1. Key Generation (Public Key Derivation)
    // ======================================================================
    let mut keygen_seeds = vec![[0u8; 160]; 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;
            // Deriving the composite public key directly from the master seed
            let binding = SWing::from_seed(black_box(seed)).unwrap();
            let vk = binding.get_pub_key();
            black_box(vk);
        })
    });

    // ======================================================================
    // 2. Signing
    // ======================================================================
    let mut sign_seed = [0u8; 160];
    getrandom::fill(&mut sign_seed).unwrap();
    
    let mut random_seed = [0u8; 64];
    getrandom::fill(&mut random_seed).unwrap();

    group.bench_function("Sign", |b| {
        b.iter(|| {
            let binding = SWing::from_seed(black_box(&sign_seed)).unwrap();
            let sig = binding.sign(
                black_box(TEST_MESSAGE),
                black_box(TEST_CONTEXT),
                black_box(random_seed),
            ).unwrap();
            black_box(sig);
        })
    });

    // ======================================================================
    // 3. Verification
    // ======================================================================
    let mut verify_seed = [0u8; 160];
    getrandom::fill(&mut verify_seed).unwrap();
    
    let mut v_random_seed = [0u8; 64];
    getrandom::fill(&mut v_random_seed).unwrap();

    // Generate the verification key and signature once for the verify loop
    let binding = SWing::from_seed(&verify_seed).unwrap();
    let vk_verify = binding.get_pub_key();
    let signature = binding.sign(TEST_MESSAGE, TEST_CONTEXT, v_random_seed).unwrap();

    group.bench_function("Verify", |b| {
        b.iter(|| {
            let res = SWing::verify(
                black_box(&vk_verify),
                black_box(TEST_MESSAGE),
                black_box(TEST_CONTEXT),
                black_box(&signature),
            ).unwrap();
            black_box(res);
        })
    });

    group.finish();
}

fn bench_pure_ed25519(c: &mut Criterion) {
    let mut group = c.benchmark_group("Baseline_Ed25519");
    group.measurement_time(Duration::from_secs(5));
    group.warm_up_time(Duration::from_secs(3));

    // ======================================================================
    // Pure Ed25519 Round Trip (KeyGen + Sign + Verify)
    // ======================================================================
    group.bench_function("Round_Trip_Total", |b| {
        let mut seed = [0u8; 32];
        getrandom::fill(&mut seed).unwrap();

        b.iter(|| {
            // 1. KeyGen
            let signing_key = SigningKey::from_bytes(black_box(&seed));
            let verifying_key = signing_key.verifying_key();

            // 2. Sign
            let signature = signing_key.sign(black_box(TEST_MESSAGE));

            // 3. Verify
            let is_valid = verifying_key.verify(black_box(TEST_MESSAGE), &signature).is_ok();

            black_box(is_valid);
        })
    });

    group.finish();
}

criterion_group!(benches, bench_s_wing, bench_pure_ed25519);
criterion_main!(benches);
