//! Benchmarks for timeboxed operations.
//!
//! This module contains performance benchmarks for time-sliced operations
//! that can be paused and resumed, enabling asynchronous computation.
//! Benchmarks cover both multiplication and modular exponentiation with
//! various time limits and input sizes.

use clock_bigint::{BigInt, MontgomeryContext, PartialResult, modexp_timeboxed, mul_timeboxed};
use criterion::{Criterion, criterion_group, criterion_main};

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

    // Benchmark timeboxed multiplication with different cycle limits
    group.bench_function("mul_timeboxed_small", |bencher| {
        let a = BigInt::from_u64(12345, 8);
        let b = BigInt::from_u64(67890, 8);

        bencher.iter(|| {
            let result = mul_timeboxed(&a, &b, 1000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => panic!("Should complete with sufficient cycles"),
            }
        });
    });

    group.bench_function("mul_timeboxed_medium", |bencher| {
        let a = BigInt::from_limbs(&[0x123456789ABCDEF0, 0xFEDCBA9876543210], 8).unwrap();
        let b = BigInt::from_limbs(&[0x1111111111111111, 0x2222222222222222], 8).unwrap();

        bencher.iter(|| {
            let result = mul_timeboxed(&a, &b, 10000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => panic!("Should complete with sufficient cycles"),
            }
        });
    });

    // Benchmark timeboxed multiplication with tight cycle limits (may not complete)
    group.bench_function("mul_timeboxed_tight_cycles", |bencher| {
        let a = BigInt::from_limbs(
            &[0x123456789ABCDEF0, 0xFEDCBA9876543210, 0x1111111111111111],
            8,
        )
        .unwrap();
        let b = BigInt::from_limbs(
            &[0x2222222222222222, 0x3333333333333333, 0x4444444444444444],
            8,
        )
        .unwrap();

        bencher.iter(|| {
            let result = mul_timeboxed(&a, &b, 10).unwrap(); // Very tight cycle limit
            // May be partial or complete depending on implementation
            let _ = result;
        });
    });

    group.finish();
}

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

    // Create test modulus and context
    let modulus = BigInt::from_limbs(
        &[
            0xFFFFFFFFFFFFFF43, // Make sure it's odd by setting LSB
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
        ],
        8,
    )
    .unwrap();
    let ctx = MontgomeryContext::new(modulus).unwrap();

    // Benchmark modular exponentiation with small exponent
    group.bench_function("modexp_timeboxed_small_exp", |bencher| {
        let base = BigInt::from_u64(12345, 8);
        let exponent = BigInt::from_u64(1000, 8); // Small exponent

        bencher.iter(|| {
            let result = modexp_timeboxed(&ctx, &base, &exponent, 10000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => panic!("Should complete with sufficient cycles"),
            }
        });
    });

    // Benchmark modular exponentiation with medium exponent
    group.bench_function("modexp_timeboxed_medium_exp", |bencher| {
        let base = BigInt::from_u64(12345, 8);
        let exponent = BigInt::from_limbs(&[0x1000, 0x0, 0x0, 0x0], 8).unwrap(); // 16384 bits

        bencher.iter(|| {
            let result = modexp_timeboxed(&ctx, &base, &exponent, 50000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => panic!("Should complete with sufficient cycles"),
            }
        });
    });

    // Benchmark modular exponentiation with tight cycle limits
    group.bench_function("modexp_timeboxed_tight_cycles", |bencher| {
        let base = BigInt::from_u64(12345, 8);
        let exponent = BigInt::from_limbs(&[0x1000, 0x0], 8).unwrap(); // Medium exponent

        bencher.iter(|| {
            let result = modexp_timeboxed(&ctx, &base, &exponent, 50).unwrap(); // Very tight cycles
            // May be partial or complete
            let _ = result;
        });
    });

    group.finish();
}

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

    // Create test modulus and context
    let modulus = BigInt::from_limbs(
        &[
            0xFFFFFFFFFFFFFF43,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
        ],
        8,
    )
    .unwrap();
    let ctx = MontgomeryContext::new(modulus).unwrap();

    // Benchmark multiplication with zero inputs
    group.bench_function("mul_timeboxed_zero_a", |bencher| {
        let a = BigInt::from_u64(0, 8);
        let b = BigInt::from_u64(12345, 8);

        bencher.iter(|| {
            let result = mul_timeboxed(&a, &b, 1000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => {
                    panic!("Zero multiplication should complete immediately")
                }
            }
        });
    });

    group.bench_function("mul_timeboxed_zero_b", |bencher| {
        let a = BigInt::from_u64(12345, 8);
        let b = BigInt::from_u64(0, 8);

        bencher.iter(|| {
            let result = mul_timeboxed(&a, &b, 1000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => {
                    panic!("Zero multiplication should complete immediately")
                }
            }
        });
    });

    // Benchmark modular exponentiation with exponent = 0
    group.bench_function("modexp_timeboxed_exp_zero", |bencher| {
        let base = BigInt::from_u64(12345, 8);
        let exponent = BigInt::from_u64(0, 8);

        bencher.iter(|| {
            let result = modexp_timeboxed(&ctx, &base, &exponent, 1000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => panic!("Exponent=0 should complete immediately"),
            }
        });
    });

    // Benchmark modular exponentiation with base = 0
    group.bench_function("modexp_timeboxed_base_zero", |bencher| {
        let base = BigInt::from_u64(0, 8);
        let exponent = BigInt::from_u64(1000, 8);

        bencher.iter(|| {
            let result = modexp_timeboxed(&ctx, &base, &exponent, 1000).unwrap();
            match result {
                PartialResult::Complete(_) => {}
                PartialResult::Partial(_) => panic!("Base=0 should complete quickly"),
            }
        });
    });

    group.finish();
}

criterion_group!(
    benches,
    bench_mul_timeboxed,
    bench_modexp_timeboxed,
    bench_timeboxed_edge_cases
);
criterion_main!(benches);
