//! Criterion micro-benchmarks for the `aura-crypto` primitives.
//!
//! Run with `cargo bench -p aura-crypto`. Benches are compile-checked in CI via
//! `cargo bench --no-run -p aura-crypto`.

use std::hint::black_box;

use criterion::{criterion_group, criterion_main, Criterion, Throughput};

use aura_crypto::kem::{kyber, x25519};
use aura_crypto::{derive_session_keys, AeadSession, HybridPrivateKey};

fn bench_kyber768_keygen(c: &mut Criterion) {
    c.bench_function("kyber768_keygen", |b| {
        b.iter(|| black_box(kyber::generate()));
    });
}

fn bench_kyber768_encaps(c: &mut Criterion) {
    let kp = kyber::generate();
    c.bench_function("kyber768_encaps", |b| {
        b.iter(|| black_box(kyber::encapsulate(black_box(&kp.ek)).expect("encaps")));
    });
}

fn bench_kyber768_decaps(c: &mut Criterion) {
    let kp = kyber::generate();
    let (ct, _ss) = kyber::encapsulate(&kp.ek).expect("encaps");
    c.bench_function("kyber768_decaps", |b| {
        b.iter(|| {
            black_box(kyber::decapsulate(black_box(&kp.dk), black_box(&ct)).expect("decaps"))
        });
    });
}

fn bench_x25519(c: &mut Criterion) {
    // Full ephemeral-static ECDH: keygen + diffie_hellman, representative of one handshake half.
    let server_secret = x25519::generate_secret();
    let server_public = x25519::public_bytes(&server_secret);
    c.bench_function("x25519_dh", |b| {
        b.iter(|| {
            let eph = x25519::generate_secret();
            let ss = x25519::diffie_hellman(&eph, black_box(&server_public));
            black_box(ss)
        });
    });
}

fn bench_hybrid_handshake(c: &mut Criterion) {
    c.bench_function("hybrid_handshake", |b| {
        b.iter(|| {
            // Full hybrid handshake: generate + encapsulate + decapsulate + derive session keys.
            let (private, public) = HybridPrivateKey::generate();
            let (ct, ss_server) = public.encapsulate();
            let ss_client = private.decapsulate(&ct).expect("decapsulate");
            let keys = derive_session_keys(&ss_client, &[0u8; 32], &[1u8; 32]);
            black_box((ss_server, keys))
        });
    });
}

fn bench_aead_1kb(c: &mut Criterion) {
    let data = vec![0xABu8; 1024];
    let aad = b"aura";
    let mut group = c.benchmark_group("aead_seal");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.bench_function("aead_1kb", |b| {
        let mut session = AeadSession::new([0x11u8; 32]);
        b.iter(|| black_box(session.seal(black_box(&data), aad)));
    });
    group.finish();
}

fn bench_aead_64kb(c: &mut Criterion) {
    let data = vec![0xCDu8; 64 * 1024];
    let aad = b"aura";
    let mut group = c.benchmark_group("aead_seal");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.bench_function("aead_64kb", |b| {
        let mut session = AeadSession::new([0x22u8; 32]);
        b.iter(|| black_box(session.seal(black_box(&data), aad)));
    });
    group.finish();
}

criterion_group!(
    benches,
    bench_kyber768_keygen,
    bench_kyber768_encaps,
    bench_kyber768_decaps,
    bench_x25519,
    bench_hybrid_handshake,
    bench_aead_1kb,
    bench_aead_64kb,
);
criterion_main!(benches);