//! Cryptographic operation benchmarks
//!
//! Run with: cargo bench

use criterion::{black_box, criterion_group, criterion_main, Criterion, BenchmarkId};
use mwrasp_crypto::{
    keys::{hkdf, hmac},
    symmetric::aes_gcm,
    pqc::{kyber, dilithium, sphincs, SecurityLevel},
};

fn bench_hkdf(c: &mut Criterion) {
    let secret = [0u8; 32];
    let salt = [1u8; 16];
    let info = b"benchmark";

    let mut group = c.benchmark_group("HKDF");

    for length in [16, 32, 64, 128].iter() {
        group.bench_with_input(
            BenchmarkId::new("derive_key", length),
            length,
            |b, &len| {
                b.iter(|| {
                    hkdf::derive_key(
                        black_box(&secret),
                        Some(&salt),
                        black_box(info),
                        len,
                    )
                })
            },
        );
    }
    group.finish();
}

fn bench_hmac(c: &mut Criterion) {
    let key = [0u8; 32];

    let mut group = c.benchmark_group("HMAC-SHA256");

    for size in [64, 256, 1024, 4096].iter() {
        let message = vec![0u8; *size];
        group.bench_with_input(
            BenchmarkId::new("compute", size),
            &message,
            |b, msg| {
                b.iter(|| hmac::hmac_sha256(black_box(&key), black_box(msg)))
            },
        );
    }

    // Benchmark verification
    let message = vec![0u8; 256];
    let tag = hmac::hmac_sha256(&key, &message);
    group.bench_function("verify_256B", |b| {
        b.iter(|| hmac::verify_hmac(black_box(&key), black_box(&message), black_box(&tag)))
    });

    group.finish();
}

fn bench_aes_gcm(c: &mut Criterion) {
    let key = [0u8; 32];
    let nonce = [0u8; 12];
    let aad = b"additional data";

    let mut group = c.benchmark_group("AES-256-GCM");

    for size in [64, 256, 1024, 4096, 16384].iter() {
        let plaintext = vec![0u8; *size];

        group.bench_with_input(
            BenchmarkId::new("encrypt", size),
            &plaintext,
            |b, pt| {
                b.iter(|| {
                    aes_gcm::encrypt_aes_gcm(
                        black_box(&key),
                        black_box(&nonce),
                        black_box(pt),
                        black_box(aad),
                    )
                })
            },
        );

        let ciphertext = aes_gcm::encrypt_aes_gcm(&key, &nonce, &plaintext, aad).unwrap();
        group.bench_with_input(
            BenchmarkId::new("decrypt", size),
            &ciphertext,
            |b, ct| {
                b.iter(|| {
                    aes_gcm::decrypt_aes_gcm(
                        black_box(&key),
                        black_box(&nonce),
                        black_box(ct),
                        black_box(aad),
                    )
                })
            },
        );
    }

    group.finish();
}

fn bench_kyber(c: &mut Criterion) {
    let mut group = c.benchmark_group("Kyber (ML-KEM)");

    // Kyber-768 (Level 3) - recommended
    group.bench_function("keygen_768", |b| {
        b.iter(|| kyber::kyber_keygen(black_box(SecurityLevel::Level3)))
    });

    let keypair = kyber::kyber_keygen(SecurityLevel::Level3).unwrap();
    group.bench_function("encapsulate_768", |b| {
        b.iter(|| kyber::kyber_encapsulate(black_box(&keypair.public_key)))
    });

    let (ciphertext, _) = kyber::kyber_encapsulate(&keypair.public_key).unwrap();
    group.bench_function("decapsulate_768", |b| {
        b.iter(|| {
            kyber::kyber_decapsulate(
                black_box(&ciphertext),
                black_box(keypair.secret_key()),
            )
        })
    });

    group.finish();
}

fn bench_dilithium(c: &mut Criterion) {
    let mut group = c.benchmark_group("Dilithium (ML-DSA)");
    let message = b"benchmark message for signing";

    // Dilithium3 (Level 3) - recommended
    group.bench_function("keygen_3", |b| {
        b.iter(|| dilithium::dilithium_keygen(black_box(SecurityLevel::Level3)))
    });

    let keypair = dilithium::dilithium_keygen(SecurityLevel::Level3).unwrap();
    group.bench_function("sign_3", |b| {
        b.iter(|| {
            dilithium::dilithium_sign(
                black_box(keypair.secret_key()),
                black_box(message),
            )
        })
    });

    let signature = dilithium::dilithium_sign(keypair.secret_key(), message).unwrap();
    group.bench_function("verify_3", |b| {
        b.iter(|| {
            dilithium::dilithium_verify(
                black_box(&keypair.public_key),
                black_box(message),
                black_box(&signature),
            )
        })
    });

    group.finish();
}

fn bench_sphincs(c: &mut Criterion) {
    let mut group = c.benchmark_group("SPHINCS+ (SLH-DSA)");
    group.sample_size(10); // SPHINCS+ is slow, reduce sample size

    let message = b"benchmark message for signing";
    let variant = sphincs::SphincsVariant::Sha2_128f; // Fast variant

    group.bench_function("keygen_sha2_128f", |b| {
        b.iter(|| sphincs::sphincs_keygen(black_box(variant)))
    });

    let keypair = sphincs::sphincs_keygen(variant).unwrap();
    group.bench_function("sign_sha2_128f", |b| {
        b.iter(|| {
            sphincs::sphincs_sign_with_variant(
                black_box(keypair.secret_key()),
                black_box(message),
                variant,
            )
        })
    });

    let signature = sphincs::sphincs_sign_with_variant(keypair.secret_key(), message, variant).unwrap();
    group.bench_function("verify_sha2_128f", |b| {
        b.iter(|| {
            sphincs::sphincs_verify_with_variant(
                black_box(&keypair.public_key),
                black_box(message),
                black_box(&signature),
                variant,
            )
        })
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_hkdf,
    bench_hmac,
    bench_aes_gcm,
    bench_kyber,
    bench_dilithium,
    bench_sphincs,
);
criterion_main!(benches);
