//! `aura bench-crypto`: quick KEM / AEAD / handshake timings.
//!
//! This is a lightweight, dependency-free micro-benchmark (no criterion) intended for an at-a-glance
//! feel of the crypto core's cost. It times, with [`std::time::Instant`], the hybrid KEM keygen,
//! encapsulation, decapsulation, a full hybrid handshake (keygen + encaps + decaps + key
//! derivation), and AEAD seal/open over 1 KiB and 64 KiB messages, then prints a table.

use std::time::{Duration, Instant};

use aura_crypto::{derive_session_keys, AeadSession, HybridPrivateKey};

/// Number of iterations per measured operation.
const ITERS: u32 = 200;

/// Run the crypto micro-benchmarks and print a results table to stdout.
pub fn run() -> anyhow::Result<()> {
    println!("aura bench-crypto — {ITERS} iterations per op (hybrid X25519 + ML-KEM-768)\n");
    println!("{:<32} {:>12} {:>14}", "operation", "avg", "ops/sec");
    println!("{}", "-".repeat(60));

    // KEM keygen.
    let keygen = time(ITERS, || {
        let _ = HybridPrivateKey::generate();
    });
    row("KEM keygen", keygen);

    // Encapsulate (server side) against a fixed public key.
    let (_sk, pk) = HybridPrivateKey::generate();
    let encaps = time(ITERS, || {
        let _ = pk.encapsulate();
    });
    row("KEM encapsulate", encaps);

    // Decapsulate (client side) of a fixed ciphertext.
    let (sk, pk) = HybridPrivateKey::generate();
    let (ct, _ss) = pk.encapsulate();
    let decaps = time(ITERS, || {
        let _ = sk.decapsulate(&ct).expect("decapsulate");
    });
    row("KEM decapsulate", decaps);

    // Full hybrid handshake: keygen + encaps + decaps + derive both directions' keys.
    let nonce = [0u8; 32];
    let handshake = time(ITERS, || {
        let (sk, pk) = HybridPrivateKey::generate();
        let (ct, server_ss) = pk.encapsulate();
        let client_ss = sk.decapsulate(&ct).expect("decapsulate");
        let _ = derive_session_keys(&server_ss, &nonce, &nonce);
        let _ = derive_session_keys(&client_ss, &nonce, &nonce);
    });
    row("full hybrid handshake", handshake);

    // AEAD seal/open round trips at two payload sizes.
    for (label, size) in [
        ("AEAD seal+open 1KiB", 1024usize),
        ("AEAD seal+open 64KiB", 64 * 1024),
    ] {
        let plaintext = vec![0xA5u8; size];
        let aad = b"aura-bench";
        let key = [7u8; 32];
        let d = time(ITERS, || {
            let mut seal = AeadSession::new(key);
            let mut open = AeadSession::new(key);
            let ct = seal.seal(&plaintext, aad);
            let pt = open.open(&ct, aad).expect("open");
            debug_assert_eq!(pt.len(), size);
        });
        row(label, d);
    }

    println!("\n(timings are wall-clock averages on this host; not a substitute for criterion)");
    Ok(())
}

/// Time `f` over `iters` iterations and return the total elapsed duration.
fn time(iters: u32, mut f: impl FnMut()) -> Duration {
    // One warm-up iteration to avoid counting first-call setup.
    f();
    let start = Instant::now();
    for _ in 0..iters {
        f();
    }
    start.elapsed()
}

/// Print one results row given the total time for [`ITERS`] iterations.
fn row(label: &str, total: Duration) {
    let avg = total / ITERS;
    let per_sec = if avg.as_secs_f64() > 0.0 {
        1.0 / avg.as_secs_f64()
    } else {
        f64::INFINITY
    };
    println!("{label:<32} {:>12} {per_sec:>14.0}", format!("{avg:?}"));
}