#![allow(clippy::unwrap_used)]
//! Benchmark: Hybrid KEM (X25519 + ML-KEM-768) vs classical-only encryption.
//!
//! Compares keypair generation, encryption, and decryption times for the
//! hybrid post-quantum scheme against classical X25519 + ChaCha20-Poly1305.

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};

use quicprochat_core::{hybrid_encrypt, hybrid_decrypt, HybridKeypair};

// ── Classical baseline (X25519 + ChaCha20-Poly1305) ─────────────────────────

use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Key, Nonce,
};
use hkdf::Hkdf;
use rand::{rngs::OsRng, RngCore};
use sha2::Sha256;
use x25519_dalek::{EphemeralSecret, PublicKey as X25519Public, StaticSecret};

struct ClassicalKeypair {
    secret: StaticSecret,
    public: X25519Public,
}

impl ClassicalKeypair {
    fn generate() -> Self {
        let secret = StaticSecret::random_from_rng(OsRng);
        let public = X25519Public::from(&secret);
        Self { secret, public }
    }
}

fn classical_encrypt(recipient_pk: &X25519Public, plaintext: &[u8]) -> Vec<u8> {
    let eph_secret = EphemeralSecret::random_from_rng(OsRng);
    let eph_public = X25519Public::from(&eph_secret);
    let shared = eph_secret.diffie_hellman(recipient_pk);

    let hk = Hkdf::<Sha256>::new(None, shared.as_bytes());
    let mut key_bytes = [0u8; 32];
    hk.expand(b"classical-bench", &mut key_bytes).unwrap();

    let mut nonce_bytes = [0u8; 12];
    OsRng.fill_bytes(&mut nonce_bytes);

    let cipher = ChaCha20Poly1305::new(Key::from_slice(&key_bytes));
    let ct = cipher
        .encrypt(Nonce::from_slice(&nonce_bytes), plaintext)
        .unwrap();

    // Wire: eph_pk(32) || nonce(12) || ciphertext
    let mut out = Vec::with_capacity(32 + 12 + ct.len());
    out.extend_from_slice(eph_public.as_bytes());
    out.extend_from_slice(&nonce_bytes);
    out.extend_from_slice(&ct);
    out
}

fn classical_decrypt(keypair: &ClassicalKeypair, envelope: &[u8]) -> Vec<u8> {
    let eph_pk = X25519Public::from(<[u8; 32]>::try_from(&envelope[..32]).unwrap());
    let nonce_bytes: [u8; 12] = envelope[32..44].try_into().unwrap();
    let ct = &envelope[44..];

    let shared = keypair.secret.diffie_hellman(&eph_pk);

    let hk = Hkdf::<Sha256>::new(None, shared.as_bytes());
    let mut key_bytes = [0u8; 32];
    hk.expand(b"classical-bench", &mut key_bytes).unwrap();

    let cipher = ChaCha20Poly1305::new(Key::from_slice(&key_bytes));
    cipher
        .decrypt(Nonce::from_slice(&nonce_bytes), ct)
        .unwrap()
}

// ── Benchmarks ──────────────────────────────────────────────────────────────

fn bench_keygen(c: &mut Criterion) {
    let mut group = c.benchmark_group("kem_keygen");
    group.bench_function("hybrid", |b| {
        b.iter(|| black_box(HybridKeypair::generate()));
    });
    group.bench_function("classical", |b| {
        b.iter(|| black_box(ClassicalKeypair::generate()));
    });
    group.finish();
}

fn bench_encrypt(c: &mut Criterion) {
    let sizes: &[(&str, usize)] = &[("100B", 100), ("1KB", 1024), ("4KB", 4096), ("64KB", 65536)];
    let mut group = c.benchmark_group("kem_encrypt");

    let hybrid_kp = HybridKeypair::generate();
    let hybrid_pk = hybrid_kp.public_key();
    let classical_kp = ClassicalKeypair::generate();

    for (label, size) in sizes {
        let payload = vec![0xABu8; *size];

        group.bench_with_input(
            BenchmarkId::new("hybrid", label),
            &payload,
            |b, payload| {
                b.iter(|| hybrid_encrypt(&hybrid_pk, black_box(payload), b"", b"").unwrap());
            },
        );

        group.bench_with_input(
            BenchmarkId::new("classical", label),
            &payload,
            |b, payload| {
                b.iter(|| classical_encrypt(&classical_kp.public, black_box(payload)));
            },
        );
    }
    group.finish();
}

fn bench_decrypt(c: &mut Criterion) {
    let sizes: &[(&str, usize)] = &[("100B", 100), ("1KB", 1024), ("4KB", 4096), ("64KB", 65536)];
    let mut group = c.benchmark_group("kem_decrypt");

    let hybrid_kp = HybridKeypair::generate();
    let hybrid_pk = hybrid_kp.public_key();
    let classical_kp = ClassicalKeypair::generate();

    for (label, size) in sizes {
        let payload = vec![0xABu8; *size];
        let hybrid_ct = hybrid_encrypt(&hybrid_pk, &payload, b"", b"").unwrap();
        let classical_ct = classical_encrypt(&classical_kp.public, &payload);

        group.bench_with_input(
            BenchmarkId::new("hybrid", label),
            &hybrid_ct,
            |b, ct| {
                b.iter(|| hybrid_decrypt(&hybrid_kp, black_box(ct), b"", b"").unwrap());
            },
        );

        group.bench_with_input(
            BenchmarkId::new("classical", label),
            &classical_ct,
            |b, ct| {
                b.iter(|| classical_decrypt(&classical_kp, black_box(ct)));
            },
        );
    }
    group.finish();
}

criterion_group!(benches, bench_keygen, bench_encrypt, bench_decrypt);
criterion_main!(benches);