//! Benchmarks for modular exponentiation operations.
//!
//! This module contains performance benchmarks for modular exponentiation
//! using Montgomery arithmetic. Benchmarks are provided for both 256-bit
//! and 2048-bit operands with various exponent sizes.

use clock_bigint::{BigInt, MontgomeryContext, mod_pow};
use criterion::{Criterion, criterion_group, criterion_main};

fn bench_modexp_u256(c: &mut Criterion) {
    let mut group = c.benchmark_group("modexp_u256");

    // Create a 256-bit modulus (must be odd)
    let modulus = BigInt::from_limbs(
        &[
            0xFFFFFFFFFFFFFF43, // Make sure it's odd by setting LSB
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
        ],
        4,
    )
    .unwrap();

    let ctx = MontgomeryContext::new(modulus).unwrap();

    // Benchmark 256-bit modular exponentiation
    group.bench_function("modexp_u256", |b| {
        let base = BigInt::from_u64(0x123456789ABCDEF0, 4);
        let exponent = BigInt::from_limbs(
            &[
                0x10001, // Small exponent for RSA-like operations
                0, 0, 0,
            ],
            4,
        )
        .unwrap();

        b.iter(|| {
            std::hint::black_box(mod_pow(&ctx, &base, &exponent)).unwrap();
        });
    });

    // Benchmark with larger exponent
    group.bench_function("modexp_u256_large_exp", |b| {
        let base = BigInt::from_u64(0x123456789ABCDEF0, 4);
        let exponent = BigInt::from_limbs(
            &[
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
            ],
            4,
        )
        .unwrap();

        b.iter(|| {
            std::hint::black_box(mod_pow(&ctx, &base, &exponent)).unwrap();
        });
    });

    group.finish();
}

fn bench_modexp_u2048(c: &mut Criterion) {
    let mut group = c.benchmark_group("modexp_u2048");

    // Create a 2048-bit modulus (must be odd)
    let mut modulus_limbs = [0u64; 32];
    modulus_limbs[0] = 0xFFFFFFFFFFFFFF43; // Make sure it's odd
    for i in 1..32 {
        modulus_limbs[i] = 0xFFFFFFFFFFFFFFFF;
    }
    let modulus = BigInt::from_limbs(&modulus_limbs, 32).unwrap();

    let ctx = MontgomeryContext::new(modulus).unwrap();

    // Benchmark 2048-bit modular exponentiation with small exponent (RSA verification)
    group.bench_function("modexp_u2048_small_exp", |b| {
        let base = BigInt::from_limbs(&[0x123456789ABCDEF0; 32], 32).unwrap();
        let exponent = BigInt::from_limbs(
            &[
                0x10001, // RSA exponent
                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0,
            ],
            32,
        )
        .unwrap();

        b.iter(|| {
            std::hint::black_box(mod_pow(&ctx, &base, &exponent)).unwrap();
        });
    });

    // Benchmark 2048-bit modular exponentiation with large exponent (key exchange)
    group.bench_function("modexp_u2048_large_exp", |b| {
        let base = BigInt::from_limbs(&[0x123456789ABCDEF0; 32], 32).unwrap();
        let exponent = BigInt::from_limbs(
            &[
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x123456789ABCDEF0,
            ],
            32,
        )
        .unwrap();

        b.iter(|| {
            std::hint::black_box(mod_pow(&ctx, &base, &exponent)).unwrap();
        });
    });

    group.finish();
}

fn bench_modexp_edge_cases(c: &mut Criterion) {
    let mut group = c.benchmark_group("modexp_edge_cases");

    // Create a 256-bit modulus
    let modulus = BigInt::from_limbs(
        &[
            0xFFFFFFFFFFFFFF43,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
        ],
        4,
    )
    .unwrap();

    let ctx = MontgomeryContext::new(modulus).unwrap();

    // Benchmark exponent = 0 (should return 1)
    group.bench_function("modexp_exp_zero", |b| {
        let base = BigInt::from_u64(0x123456789ABCDEF0, 4);
        let exponent = BigInt::from_u64(0, 4);

        b.iter(|| {
            std::hint::black_box(mod_pow(&ctx, &base, &exponent)).unwrap();
        });
    });

    // Benchmark base = 0 (should return 0 for exp > 0)
    group.bench_function("modexp_base_zero", |b| {
        let base = BigInt::from_u64(0, 4);
        let exponent = BigInt::from_u64(0x123456789ABCDEF0, 4);

        b.iter(|| {
            std::hint::black_box(mod_pow(&ctx, &base, &exponent)).unwrap();
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_modexp_u256,
    bench_modexp_u2048,
    bench_modexp_edge_cases
);
criterion_main!(benches);
