//! Throughput of the public batch-verification API across batch sizes. Each signature is from an
//! independent key over an independent 32-byte message (no key/message reuse to amortize), the
//! harder case for the underlying MSM. Fixture generation is not timed.

use commonware_allocator::mimalloc;
use commonware_cryptography::{BatchEntry, BatchVerifier as _};
use commonware_cryptography_curve25519::signing::{Signature, SigningKey, VerifyingKey};
use commonware_math::algebra::Random;
use commonware_parallel::{Rayon, Sequential, Strategy};
use commonware_utils::{NZUsize, TestRng, test_rng};
use criterion::{BatchSize, Criterion, Throughput, criterion_group, criterion_main};
use std::hint::black_box;

#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc::new(mimalloc::Config::DEFAULT);

const NAMESPACE: &[u8] = b"_COMMONWARE_CRYPTOGRAPHY_CURVE25519_BATCH_VERIFY_BENCH";
const MESSAGE_SIZE: usize = 32;

/// Generates `n` valid `(verifying key, signature, message)` triples, signed by independent keys
/// over independent 32-byte messages using the public API.
fn generate_batch(n: usize) -> Vec<(VerifyingKey, Signature, Vec<u8>)> {
    let mut rng = test_rng();
    (0..n)
        .map(|i| {
            let signing_key = SigningKey::random(&mut rng);

            let mut message = vec![0u8; MESSAGE_SIZE];
            message[..8].copy_from_slice(&(i as u64).to_le_bytes());
            let signature = signing_key.sign(NAMESPACE, &message);

            (signing_key.verifying_key(), signature, message)
        })
        .collect()
}

fn verify_batch(
    rng: &mut TestRng,
    batch: &[(VerifyingKey, Signature, Vec<u8>)],
    strategy: &impl Strategy,
) -> bool {
    VerifyingKey::verify_batch(
        rng,
        batch,
        |_, (verifying_key, signature, message)| BatchEntry {
            namespace: NAMESPACE,
            message,
            public_key: verifying_key,
            signature,
        },
        strategy,
    )
}

/// Nested so `module_path!()` includes the crate::module separator the benchmark-name lint (and
/// the benchmark-tracking dashboard) expect.
mod verify_batch_bytes_bench {
    use super::{
        BatchSize, Criterion, NZUsize, Rayon, Sequential, TestRng, Throughput, black_box,
        generate_batch, verify_batch,
    };

    pub fn bench(c: &mut Criterion) {
        let mut group = c.benchmark_group(module_path!());
        for n in [1, 10, 100, 1_000, 10_000, 16_384] {
            let batch = generate_batch(n);
            for concurrency in [1, 8, 16, 32] {
                let rayon = (concurrency > 1).then(|| Rayon::new(NZUsize!(concurrency)).unwrap());
                group.throughput(Throughput::Elements(n as u64));
                group.bench_function(format!("sigs={n} conc={concurrency}"), |b| {
                    b.iter_batched(
                        || TestRng::new(1),
                        |mut rng| {
                            if let Some(rayon) = rayon.as_ref() {
                                black_box(verify_batch(&mut rng, &batch, rayon))
                            } else {
                                black_box(verify_batch(&mut rng, &batch, &Sequential))
                            }
                        },
                        BatchSize::SmallInput,
                    );
                });
            }
        }
    }
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = verify_batch_bytes_bench::bench
}

criterion_main!(benches);