//! Modular exponentiation using Montgomery ladder.
//!
//! Implements constant-time modular exponentiation using the Montgomery ladder
//! algorithm with constant-time conditional swaps.

use crate::ct;
use crate::error::Result;
use crate::montgomery::{MontgomeryContext, from_mont, mont_mul, to_mont};
use crate::types::BigIntCore;

#[cfg(feature = "alloc")]
use crate::types::BigInt;

/// Modular exponentiation: base^exponent mod modulus
///
/// Uses Montgomery ladder algorithm for constant-time execution.
#[cfg(feature = "alloc")]
pub fn mod_pow(ctx: &MontgomeryContext, base: &BigInt, exponent: &BigInt) -> Result<BigInt> {
    // Handle special cases
    if exponent.is_zero() {
        // base^0 = 1 mod m
        let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
        return Ok(one);
    }

    if base.is_zero() {
        // 0^exponent = 0 mod m (if exponent > 0)
        return Ok(BigInt::from_u64(0, ctx.modulus().max_limbs()));
    }

    // Convert base to Montgomery form
    let x_mont = to_mont(ctx, base)?;

    // Convert 1 to Montgomery form (R̃ = to_mont(1))
    let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
    let r_mont = to_mont(ctx, &one)?;

    // Initialize Montgomery ladder: R0 = R̃, R1 = x̃
    let mut r0 = r_mont;
    let mut r1 = x_mont;

    // Get exponent bits (from MSB to LSB)
    let exp_limbs = exponent.limbs();
    let mut bit_pos = 0;

    // Find the most significant bit
    let mut msb_found = false;
    for i in (0..exp_limbs.len()).rev() {
        if exp_limbs[i] != 0 {
            bit_pos = (i * 64) + (63 - exp_limbs[i].leading_zeros() as usize);
            msb_found = true;
            break;
        }
    }

    if !msb_found {
        // Exponent is zero (should have been caught above, but handle it)
        let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
        return Ok(one);
    }

    // Process bits from MSB to LSB
    for i in (0..=bit_pos).rev() {
        let limb_idx = i / 64;
        let bit_idx = i % 64;
        let bit = if limb_idx < exp_limbs.len() {
            (exp_limbs[limb_idx] >> bit_idx) & 1
        } else {
            0
        };

        // Constant-time conditional swap
        swap_bigint(bit, &mut r0, &mut r1);

        // Montgomery ladder step
        r1 = mont_mul(ctx, &r0, &r1)?;
        r0 = mont_mul(ctx, &r0, &r0)?;

        // Swap again
        swap_bigint(bit, &mut r0, &mut r1);
    }

    // Convert result from Montgomery form
    from_mont(ctx, &r0)
}

/// Constant-time swap of two BigInt values.
fn swap_bigint(choice: u64, a: &mut BigInt, b: &mut BigInt) {
    let a_limbs = a.limbs_mut();
    let b_limbs = b.limbs_mut();
    let len = a_limbs.len().min(b_limbs.len());

    for i in 0..len {
        ct::ct_swap(choice, &mut a_limbs[i], &mut b_limbs[i]);
    }

    // Also swap signs
    let a_sign = a.sign();
    let b_sign = b.sign();
    let new_a_sign = ct::ct_select(choice, b_sign as u64, a_sign as u64) != 0;
    let new_b_sign = ct::ct_select(choice, a_sign as u64, b_sign as u64) != 0;
    a.set_sign(new_a_sign);
    b.set_sign(new_b_sign);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mod_pow_exponent_zero() {
        let modulus = BigInt::from_u64(97, 10);
        let ctx = MontgomeryContext::new(modulus).unwrap();
        let base = BigInt::from_u64(5, 10);
        let exponent = BigInt::from_u64(0, 10);
        let result = mod_pow(&ctx, &base, &exponent).unwrap();
        assert_eq!(result.limbs()[0], 1);
    }
}
