//! Benchmarks comparing Ring-LPR OPRF (OT-based) vs Fast OPRF (OT-free)
//!
//! Run with: cargo bench

use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use rand::SeedableRng;
use rand_chacha::ChaCha20Rng;

use opaque_lattice::oprf::fast_oprf::{
    PublicParams, ServerKey, client_blind as fast_client_blind, client_finalize as fast_finalize,
    evaluate as fast_evaluate, server_evaluate as fast_server_evaluate,
};
use opaque_lattice::oprf::ring_lpr::{
    RingLprKey, client_blind as lpr_client_blind, client_finalize as lpr_finalize,
    server_evaluate as lpr_server_evaluate,
};
use opaque_lattice::oprf::unlinkable_oprf::{
    UnlinkablePublicParams, UnlinkableServerKey, client_blind_unlinkable,
    client_finalize_unlinkable, evaluate_unlinkable, server_evaluate_unlinkable,
};

/// Benchmark Fast OPRF (OT-free) - full protocol
fn bench_fast_oprf(c: &mut Criterion) {
    let mut group = c.benchmark_group("fast_oprf");

    let pp = PublicParams::generate(b"benchmark-params");
    let key = ServerKey::generate(&pp, b"benchmark-key");
    let password = b"benchmark-password-12345";

    // Benchmark client blind
    group.bench_function("client_blind", |b| {
        b.iter(|| fast_client_blind(&pp, password))
    });

    // Benchmark server evaluate
    let (state, blinded) = fast_client_blind(&pp, password);
    group.bench_function("server_evaluate", |b| {
        b.iter(|| fast_server_evaluate(&key, &blinded))
    });

    // Benchmark client finalize
    let response = fast_server_evaluate(&key, &blinded);
    group.bench_function("client_finalize", |b| {
        let state = state.clone();
        b.iter(|| fast_finalize(&state, key.public_key(), &response))
    });

    // Benchmark full protocol
    group.bench_function("full_protocol", |b| {
        b.iter(|| fast_evaluate(&pp, &key, password))
    });

    group.finish();
}

/// Benchmark Ring-LPR OPRF (OT-based) - full protocol
fn bench_ring_lpr_oprf(c: &mut Criterion) {
    let mut group = c.benchmark_group("ring_lpr_oprf");

    let mut rng = ChaCha20Rng::seed_from_u64(12345);
    let key = RingLprKey::generate(&mut rng);
    let password = b"benchmark-password-12345";

    // Benchmark client blind
    group.bench_function("client_blind", |b| {
        let mut rng = ChaCha20Rng::seed_from_u64(99999);
        b.iter(|| lpr_client_blind(&mut rng, password))
    });

    // Benchmark server evaluate
    let mut rng2 = ChaCha20Rng::seed_from_u64(88888);
    let (state, blinded) = lpr_client_blind(&mut rng2, password).unwrap();
    group.bench_function("server_evaluate", |b| {
        b.iter(|| lpr_server_evaluate(&key, &blinded))
    });

    // Benchmark client finalize
    let evaluated = lpr_server_evaluate(&key, &blinded).unwrap();
    group.bench_function("client_finalize", |b| {
        let mut rng = ChaCha20Rng::seed_from_u64(77777);
        let (state, _) = lpr_client_blind(&mut rng, password).unwrap();
        b.iter(|| {
            // Need to re-create state each time since finalize consumes it
            let mut rng = ChaCha20Rng::seed_from_u64(77777);
            let (state, _) = lpr_client_blind(&mut rng, password).unwrap();
            lpr_finalize(state, &evaluated)
        })
    });

    // Benchmark full protocol
    group.bench_function("full_protocol", |b| {
        let mut rng = ChaCha20Rng::seed_from_u64(66666);
        b.iter(|| {
            let (state, blinded) = lpr_client_blind(&mut rng, password).unwrap();
            let evaluated = lpr_server_evaluate(&key, &blinded).unwrap();
            lpr_finalize(state, &evaluated)
        })
    });

    group.finish();
}

/// Compare all three protocols side-by-side
fn bench_comparison(c: &mut Criterion) {
    let mut group = c.benchmark_group("oprf_comparison");

    let pp = PublicParams::generate(b"benchmark-params");
    let fast_key = ServerKey::generate(&pp, b"benchmark-key");

    let unlink_pp = UnlinkablePublicParams::generate(b"benchmark-params");
    let unlink_key = UnlinkableServerKey::generate(&unlink_pp, b"benchmark-key");

    let mut rng = ChaCha20Rng::seed_from_u64(12345);
    let lpr_key = RingLprKey::generate(&mut rng);

    let passwords = [
        b"short".as_slice(),
        b"medium-password-123".as_slice(),
        b"this-is-a-very-long-password-that-tests-longer-inputs".as_slice(),
    ];

    for password in &passwords {
        let len = password.len();

        group.bench_with_input(BenchmarkId::new("fast_oprf", len), password, |b, pwd| {
            b.iter(|| fast_evaluate(&pp, &fast_key, pwd))
        });

        group.bench_with_input(
            BenchmarkId::new("unlinkable_oprf", len),
            password,
            |b, pwd| b.iter(|| evaluate_unlinkable(&unlink_pp, &unlink_key, pwd)),
        );

        group.bench_with_input(
            BenchmarkId::new("ring_lpr_oprf", len),
            password,
            |b, pwd| {
                let mut rng = ChaCha20Rng::seed_from_u64(55555);
                b.iter(|| {
                    let (state, blinded) = lpr_client_blind(&mut rng, pwd).unwrap();
                    let evaluated = lpr_server_evaluate(&lpr_key, &blinded).unwrap();
                    lpr_finalize(state, &evaluated)
                })
            },
        );
    }

    group.finish();
}

/// Benchmark Unlinkable OPRF - full protocol
fn bench_unlinkable_oprf(c: &mut Criterion) {
    let mut group = c.benchmark_group("unlinkable_oprf");

    let pp = UnlinkablePublicParams::generate(b"benchmark-params");
    let key = UnlinkableServerKey::generate(&pp, b"benchmark-key");
    let password = b"benchmark-password-12345";

    group.bench_function("client_blind", |b| {
        b.iter(|| client_blind_unlinkable(&pp, password))
    });

    let (state, blinded) = client_blind_unlinkable(&pp, password);
    group.bench_function("server_evaluate", |b| {
        b.iter(|| server_evaluate_unlinkable(&key, &blinded))
    });

    let response = server_evaluate_unlinkable(&key, &blinded);
    group.bench_function("client_finalize", |b| {
        let state = state.clone();
        b.iter(|| client_finalize_unlinkable(&state, key.public_key(), &response))
    });

    group.bench_function("full_protocol", |b| {
        b.iter(|| evaluate_unlinkable(&pp, &key, password))
    });

    group.finish();
}

/// Benchmark message sizes
fn bench_message_sizes(c: &mut Criterion) {
    println!("\n=== Message Size Comparison ===\n");

    // Fast OPRF messages
    let pp = PublicParams::generate(b"benchmark-params");
    let fast_key = ServerKey::generate(&pp, b"benchmark-key");
    let (_, blinded) = fast_client_blind(&pp, b"password");
    let response = fast_server_evaluate(&fast_key, &blinded);

    println!("Fast OPRF:");
    println!("  Client -> Server (BlindedInput): ~{} bytes", 256 * 4); // RingElement
    println!("  Server -> Client (Response): ~{} bytes", 256 * 4 + 256); // RingElement + helper

    // Ring-LPR OPRF messages
    let mut rng = ChaCha20Rng::seed_from_u64(12345);
    let lpr_key = RingLprKey::generate(&mut rng);
    let (_, lpr_blinded) = lpr_client_blind(&mut rng, b"password").unwrap();
    let lpr_evaluated = lpr_server_evaluate(&lpr_key, &lpr_blinded).unwrap();

    let lpr_blind_size = lpr_blinded.to_bytes().len();
    let lpr_eval_size = lpr_evaluated.to_bytes().len();

    println!("\nRing-LPR OPRF:");
    println!(
        "  Client -> Server (BlindedInput): {} bytes",
        lpr_blind_size
    );
    println!(
        "  Server -> Client (EvaluatedOutput): {} bytes",
        lpr_eval_size
    );

    println!(
        "\nSpeedup factor (message size): {:.1}x",
        lpr_blind_size as f64 / (256.0 * 4.0)
    );
    println!();
}

criterion_group!(
    benches,
    bench_fast_oprf,
    bench_unlinkable_oprf,
    bench_ring_lpr_oprf,
    bench_comparison,
);

criterion_main!(benches);