//! Montgomery modular arithmetic.
//!
//! Implements Montgomery reduction and multiplication for efficient
//! modular arithmetic operations, especially for modular exponentiation.

use crate::ct;
use crate::error::{BigIntError, Result};
use crate::limbs::{Limb, limb_add_carry, limb_mul_wide, limb_sub_borrow};
use crate::types::BigIntCore;

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

#[cfg(feature = "alloc")]
/// Montgomery context containing precomputed parameters.
///
/// Stores:
/// - R mod m
/// - R² mod m
/// - m' = -m⁻¹ mod B
/// - Modulus m
#[derive(Clone, Debug)]
pub struct MontgomeryContext {
    /// Modulus m (must be odd)
    modulus: BigInt,
    /// R mod m where R = B^n
    r_mod_m: BigInt,
    /// R² mod m where R = B^n
    r2_mod_m: BigInt,
    /// m' = -m⁻¹ mod B (64-bit value)
    m_prime: Limb,
    /// Number of limbs in modulus
    n: usize,
}

#[cfg(feature = "alloc")]
impl MontgomeryContext {
    /// Create a new Montgomery context for the given modulus.
    ///
    /// The modulus MUST be odd and non-zero.
    pub fn new(modulus: BigInt) -> Result<Self> {
        // Validate modulus
        if modulus.is_zero() {
            return Err(BigIntError::InvalidModulus);
        }

        // Check if modulus is odd (lowest bit must be 1)
        let modulus_limbs = modulus.limbs();
        if modulus_limbs.is_empty() || (modulus_limbs[0] & 1) == 0 {
            return Err(BigIntError::InvalidModulus);
        }

        let n = modulus.limb_count();
        let _max_limbs = modulus.max_limbs();

        // Compute m' = -m⁻¹ mod B
        // Using extended Euclidean algorithm for 64-bit
        let m0 = modulus_limbs[0];
        let m_prime = compute_m_prime(m0);

        // Compute R = 2^(64*n) mod m
        let r_mod_m = compute_r_mod_m(&modulus, n)?;

        // Compute R² mod m = (R * R) % modulus
        let r2_mod_m = compute_r2_mod_m(&modulus, &r_mod_m)?;

        Ok(Self {
            modulus,
            r_mod_m,
            r2_mod_m,
            m_prime,
            n,
        })
    }

    /// Get the modulus.
    pub fn modulus(&self) -> &BigInt {
        &self.modulus
    }

    /// Get R (2^(64*n) mod modulus).
    pub fn r(&self) -> &BigInt {
        &self.r_mod_m
    }

    /// Get R² (R² mod modulus).
    pub fn r_squared(&self) -> &BigInt {
        &self.r2_mod_m
    }

    /// Get the number of limbs.
    pub fn limb_count(&self) -> usize {
        self.n
    }

    /// Get m' (-m⁻¹ mod 2^64).
    pub fn m_prime(&self) -> Limb {
        self.m_prime
    }
}

/// Constant-time comparison: returns true if a >= b
fn bigint_ge(a: &BigInt, b: &BigInt) -> bool {
    let a_limbs = a.limbs();
    let b_limbs = b.limbs();

    // Handle signs first
    let a_sign = a.sign();
    let b_sign = b.sign();

    // If signs differ: negative < positive, so if a is negative and b is positive, a < b
    if a_sign != b_sign {
        return !a_sign; // a >= b if a is not negative (i.e., a is positive)
    }

    // Same signs: compare magnitudes
    let mut result = 0u64; // 1 if a >= b, 0 otherwise

    // Compare from MSB to LSB
    let max_len = a_limbs.len().max(b_limbs.len());
    for i in (0..max_len).rev() {
        let a_val = if i < a_limbs.len() { a_limbs[i] } else { 0 };
        let b_val = if i < b_limbs.len() { b_limbs[i] } else { 0 };

        let gt = ct::ct_gt(a_val, b_val);
        let eq = ct::ct_eq(a_val, b_val);

        // If we found a difference, this determines the result
        result = ct::ct_select(gt, 1, result);
        result = ct::ct_select(eq, result, gt); // If equal so far, keep previous result
    }

    // For negative numbers, the comparison is reversed
    if a_sign {
        result = 1 - result;
    }

    result != 0
}

/// Compute m' = -m⁻¹ mod 2^64.
///
/// Uses the extended Euclidean algorithm to find the modular inverse.
/// Since m is odd (as required by Montgomery), it is guaranteed to have an inverse mod 2^64.
fn compute_m_prime(m: Limb) -> Limb {
    // We need to solve: m * x ≡ 1 mod 2^64
    // Then m' = -x mod 2^64
    //
    // Using extended Euclidean algorithm adapted for 64-bit arithmetic

    // Since we're working mod 2^64 and m is odd, we can use a simplified approach
    // The inverse can be computed using the formula for modular inverse in binary fields
    // For small m (like in tests), we can use a simple search

    if m == 1 {
        return 0u64.wrapping_sub(1); // -1 mod 2^64 = 2^64 - 1
    }

    // For general case, use extended Euclidean algorithm with u128 to avoid overflow
    let modulus = 1u128 << 64;
    let mut a = m as u128;
    let mut b = modulus;
    let mut x = 0i128;
    let mut y = 1i128;
    let mut last_x = 1i128;
    let mut last_y = 0i128;

    while b != 0 {
        let quotient = (a / b) as i128;
        let temp_b = b;
        b = a % b;
        a = temp_b;

        let temp_x = x;
        x = last_x - quotient * x;
        last_x = temp_x;

        let temp_y = y;
        y = last_y - quotient * y;
        last_y = temp_y;
    }

    // last_x contains the inverse of m mod 2^64
    let inverse = last_x as u64;

    // m' = -inverse mod 2^64
    0u64.wrapping_sub(inverse)
}

/// Compute R = 2^(64*n) mod m
fn compute_r_mod_m(modulus: &BigInt, n: usize) -> Result<BigInt> {
    // R = 2^(64*n) mod m
    // For single limb, use fast computation
    if n == 1 {
        let m = modulus.limbs()[0];
        // Compute 2^64 mod m using fast exponentiation
        // 2^64 mod m = (2^32 mod m)^2 mod m, but even better:
        // We can compute it directly since m is small
        let mut result = 1u128;
        for _ in 0..64 {
            result = (result * 2) % (m as u128);
        }
        return Ok(BigInt::from_u64(result as u64, modulus.max_limbs()));
    }

    // For multi-limb, we need to compute 2^(64*n) mod m
    // This is equivalent to left-shifting 1 by (64*n) bits, then taking mod m
    // We can use repeated squaring or just compute it step by step
    let mut result = BigInt::from_u64(1, modulus.max_limbs());

    // Shift left by 64*n bits, which is equivalent to multiplying by 2^(64*n)
    // We do this in steps to avoid creating huge intermediate values
    for _ in 0..(64 * n) {
        result = crate::add::add(&result, &result)?; // Multiply by 2
        // Take modulo at each step to keep the number small
        if bigint_ge(&result, modulus) {
            result = crate::sub::sub(&result, modulus)?;
        }
    }

    Ok(result)
}

/// Compute R² = (R * R) mod m
fn compute_r2_mod_m(modulus: &BigInt, r: &BigInt) -> Result<BigInt> {
    // R² = (R * R) mod m
    if modulus.limb_count() == 1 {
        let m = modulus.limbs()[0];
        let r_val = r.limbs()[0];
        let r2_val = ((r_val as u128) * (r_val as u128)) % (m as u128);
        Ok(BigInt::from_u64(r2_val as u64, modulus.max_limbs()))
    } else {
        // For multi-limb, use proper modular multiplication
        let r_squared = crate::mul::mul(r, r)?;
        // For now, implement a simple modulo operation for this case
        // Since r_squared < m * m, we can do repeated subtraction
        let mut result = r_squared;
        while bigint_ge(&result, modulus) {
            result = crate::sub::sub(&result, modulus)?;
        }
        Ok(result)
    }
}

/// Montgomery reduction: MontReduce(T) → t = T × R⁻¹ mod m
///
/// Where T is a 2n-limb value.
#[cfg(feature = "alloc")]
pub fn mont_reduce(ctx: &MontgomeryContext, t: &BigInt) -> Result<BigInt> {
    let n = ctx.n;
    let modulus_limbs = ctx.modulus().limbs();
    let t_limbs = t.limbs();

    // T may have fewer than 2n limbs - will be padded below

    let max_limbs = ctx.modulus().max_limbs();
    let mut result = BigInt::new(max_limbs);
    result.ensure_capacity(n + 1)?;

    // Create a working copy of T
    let mut t_work = t_limbs.to_vec();
    if t_work.len() < 2 * n {
        t_work.resize(2 * n, 0);
    }

    // Montgomery reduction algorithm
    for i in 0..n {
        let u_i = (t_work[i].wrapping_mul(ctx.m_prime)) as Limb;

        // (T[i .. i+n]) += u_i × m[0 .. n]
        let mut carry = 0u64;
        for j in 0..n {
            let wide = limb_mul_wide(u_i, modulus_limbs[j]);
            let (lo, hi) = (wide.lo, wide.hi);
            let sum_idx = i + j;
            if sum_idx < t_work.len() {
                let (sum, carry1) = limb_add_carry(t_work[sum_idx], lo, carry);
                t_work[sum_idx] = sum;
                carry = carry1 + hi;
            }
        }
        // Add carry to next limb
        if i + n < t_work.len() {
            let (sum, _) = limb_add_carry(t_work[i + n], 0, carry);
            t_work[i + n] = sum;
        }
    }

    // Extract result: t = T[n .. 2n]
    let result_limbs = result.limbs_mut();
    for i in 0..n {
        if i < result_limbs.len() {
            result_limbs[i] = t_work[n + i];
        }
    }

    // Conditional subtraction: if t ≥ m, then t -= m
    // Use constant-time masked subtraction
    let modulus_limbs = ctx.modulus().limbs();
    let mut borrow = 0u64;
    let mut needs_sub = 0u64;

    // Compare t and m (constant-time)
    for i in (0..n).rev() {
        let t_val = if i < result_limbs.len() {
            result_limbs[i]
        } else {
            0
        };
        let m_val = if i < modulus_limbs.len() {
            modulus_limbs[i]
        } else {
            0
        };
        let gt = ct::ct_gt(t_val, m_val);
        let eq = ct::ct_eq(t_val, m_val);
        needs_sub = ct::ct_select(needs_sub, gt | (eq & needs_sub), gt);
    }

    // Perform conditional subtraction
    for i in 0..n {
        let t_val = if i < result_limbs.len() {
            result_limbs[i]
        } else {
            0
        };
        let m_val = if i < modulus_limbs.len() {
            modulus_limbs[i]
        } else {
            0
        };
        let (diff, borrow_out) = limb_sub_borrow(t_val, m_val & needs_sub, borrow);
        if i < result_limbs.len() {
            result_limbs[i] = diff;
        }
        borrow = borrow_out & needs_sub;
    }

    result.canonicalize()?;
    Ok(result)
}

/// Montgomery multiplication: MontMul(ã, b̃) → c̃ = ã × b̃ × R⁻¹ mod m
#[cfg(feature = "alloc")]
pub fn mont_mul(ctx: &MontgomeryContext, a: &BigInt, b: &BigInt) -> Result<BigInt> {
    // Special case for single-limb modulus using the standard Montgomery multiplication formula
    if ctx.n == 1 {
        let m = ctx.modulus().limbs()[0];
        let a_val = a.limbs()[0];
        let b_val = b.limbs()[0];
        let m_prime = ctx.m_prime();
        let _m = m; // Suppress unused variable warning

        // Montgomery multiplication: c = (a * b + ((a * b * m') mod R) * m) / R
        // where R = 2^64
        let ab = (a_val as u128) * (b_val as u128);
        let ab_mp = ab.wrapping_mul(m_prime as u128);
        let t = ab + ((ab_mp & ((1u128 << 64) - 1)) * (m as u128));
        let result = (t >> 64) as u64;

        // Final reduction
        let final_result = if result >= m { result - m } else { result };

        return Ok(BigInt::from_u64(final_result, ctx.modulus().max_limbs()));
    }

    // First compute a × b (full multiplication)
    let ab = crate::mul::mul(a, b)?;

    // Then apply Montgomery reduction
    mont_reduce(ctx, &ab)
}

/// Convert a value to Montgomery form: to_mont(x) = MontMul(x, R mod m)
#[cfg(feature = "alloc")]
pub fn to_mont(ctx: &MontgomeryContext, x: &BigInt) -> Result<BigInt> {
    // Special case for single-limb modulus
    if ctx.n == 1 {
        let m = ctx.modulus().limbs()[0];
        let x_val = x.limbs()[0];
        let r_val = ctx.r().limbs()[0];

        // Compute (x * R) mod m directly
        let product = (x_val as u128) * (r_val as u128) % (m as u128);
        return Ok(BigInt::from_u64(product as u64, ctx.modulus().max_limbs()));
    }

    mont_mul(ctx, x, &ctx.r_mod_m)
}

/// Convert a value from Montgomery form: from_mont(x̃) = MontReduce(x̃)
#[cfg(feature = "alloc")]
pub fn from_mont(ctx: &MontgomeryContext, x: &BigInt) -> Result<BigInt> {
    // Special case for single-limb modulus
    if ctx.n == 1 {
        // For Montgomery form, from_mont(x̃) = mont_mul(x̃, 1)
        let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
        return mont_mul(ctx, x, &one);
    }

    // For from_mont, we need to treat x as a 2n-limb value
    // For now, we'll pad it
    let n = ctx.n;
    let x_limbs = x.limbs();

    let max_limbs = ctx.modulus().max_limbs();
    let mut padded = BigInt::new(max_limbs);
    padded.ensure_capacity(2 * n)?;
    let padded_limbs = padded.limbs_mut();
    padded_limbs[..x_limbs.len().min(2 * n)].copy_from_slice(&x_limbs[..x_limbs.len().min(2 * n)]);

    mont_reduce(ctx, &padded)
}

/// Montgomery addition: (ã + b̃) mod m
#[cfg(feature = "alloc")]
pub fn mont_add(ctx: &MontgomeryContext, a: &BigInt, b: &BigInt) -> Result<BigInt> {
    let sum = crate::add::add(a, b)?;

    // Conditional reduction: if sum >= m, subtract m
    let modulus = ctx.modulus();
    let sum_limbs = sum.limbs();
    let mod_limbs = modulus.limbs();

    // Compare (constant-time)
    let mut needs_reduction = 0u64;
    let n = sum_limbs.len().max(mod_limbs.len());
    for i in (0..n).rev() {
        let s_val = if i < sum_limbs.len() { sum_limbs[i] } else { 0 };
        let m_val = if i < mod_limbs.len() { mod_limbs[i] } else { 0 };
        let gt = ct::ct_gt(s_val, m_val);
        let eq = ct::ct_eq(s_val, m_val);
        needs_reduction = ct::ct_select(needs_reduction, gt | (eq & needs_reduction), gt);
    }

    // Conditional subtraction
    if needs_reduction != 0 {
        crate::sub::sub(&sum, modulus)
    } else {
        Ok(sum)
    }
}

/// Montgomery subtraction: (ã - b̃) mod m
#[cfg(feature = "alloc")]
pub fn mont_sub(ctx: &MontgomeryContext, a: &BigInt, b: &BigInt) -> Result<BigInt> {
    let diff = crate::sub::sub(a, b)?;

    // If result is negative, add modulus
    if diff.sign() {
        crate::add::add(&diff, ctx.modulus())
    } else {
        Ok(diff)
    }
}

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

    #[test]
    fn test_montgomery_context_creation() {
        let modulus = BigInt::from_u64(97, 10); // Prime number
        let ctx = MontgomeryContext::new(modulus);
        assert!(ctx.is_ok());
    }

    #[test]
    fn test_montgomery_context_even_modulus() {
        let modulus = BigInt::from_u64(96, 10); // Even number
        let ctx = MontgomeryContext::new(modulus);
        assert!(ctx.is_err());
    }
}
