//! Benchmarks for Montgomery arithmetic operations.
//!
//! This module contains performance benchmarks for Montgomery multiplication,
//! modular reduction, and conversions between standard and Montgomery forms.
//! Benchmarks are provided for both 256-bit and 2048-bit operands.

use clock_bigint::{BigInt, MontgomeryContext, from_mont, mont_mul, mont_reduce, to_mont};
use criterion::{Criterion, criterion_group, criterion_main};

fn bench_montgomery_u256(c: &mut Criterion) {
    let mut group = c.benchmark_group("montgomery_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 Montgomery multiplication
    group.bench_function("mont_mul_u256", |bencher| {
        let a = BigInt::from_u64(0x123456789ABCDEF0, 4);
        let b_val = BigInt::from_u64(0xFEDCBA9876543210, 4);

        // Convert to Montgomery form
        let a_mont = to_mont(&ctx, &a).unwrap();
        let b_mont = to_mont(&ctx, &b_val).unwrap();

        bencher.iter(|| {
            std::hint::black_box(mont_mul(&ctx, &a_mont, &b_mont)).unwrap();
        });
    });

    // Benchmark to_mont conversion
    group.bench_function("to_mont_u256", |bencher| {
        let a = BigInt::from_u64(0x123456789ABCDEF0, 4);

        bencher.iter(|| {
            std::hint::black_box(to_mont(&ctx, &a)).unwrap();
        });
    });

    // Benchmark from_mont conversion
    group.bench_function("from_mont_u256", |bencher| {
        let a = BigInt::from_u64(0x123456789ABCDEF0, 4);
        let a_mont = to_mont(&ctx, &a).unwrap();

        bencher.iter(|| {
            std::hint::black_box(from_mont(&ctx, &a_mont)).unwrap();
        });
    });

    // Benchmark mont_reduce
    group.bench_function("mont_reduce_u256", |bencher| {
        // Create a value that needs reduction (512-bit input for 256-bit modulus)
        let t = BigInt::from_limbs(
            &[
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x1111111111111111,
                0x2222222222222222,
                0x3333333333333333,
                0x4444444444444444,
                0x5555555555555555,
                0x6666666666666666,
            ],
            8,
        )
        .unwrap();

        bencher.iter(|| {
            std::hint::black_box(mont_reduce(&ctx, &t)).unwrap();
        });
    });

    group.finish();
}

fn bench_montgomery_u2048(c: &mut Criterion) {
    let mut group = c.benchmark_group("montgomery_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 Montgomery multiplication for 2048-bit
    group.bench_function("mont_mul_u2048", |bencher| {
        let a = BigInt::from_limbs(&[0x123456789ABCDEF0; 32], 32).unwrap();
        let b_val = BigInt::from_limbs(&[0xFEDCBA9876543210; 32], 32).unwrap();

        // Convert to Montgomery form
        let a_mont = to_mont(&ctx, &a).unwrap();
        let b_mont = to_mont(&ctx, &b_val).unwrap();

        bencher.iter(|| {
            std::hint::black_box(mont_mul(&ctx, &a_mont, &b_mont)).unwrap();
        });
    });

    // Benchmark to_mont conversion for 2048-bit
    group.bench_function("to_mont_u2048", |bencher| {
        let a = BigInt::from_limbs(&[0x123456789ABCDEF0; 32], 32).unwrap();

        bencher.iter(|| {
            std::hint::black_box(to_mont(&ctx, &a)).unwrap();
        });
    });

    group.finish();
}

criterion_group!(benches, bench_montgomery_u256, bench_montgomery_u2048);
criterion_main!(benches);
