//! Modular multiplication with sparse-prime folding reduction.
use super::U256;

impl U256 {
    /// Computes `(self * other) mod modulus` via widening multiplication
    /// followed by modular reduction.
    ///
    /// First computes the full 512-bit product using schoolbook long
    /// multiplication. For sparse primes where `2^256 - modulus` fits in
    /// a single u64 (e.g. secp256k1's P), uses a fast two-round folding
    /// reduction that replaces expensive multi-precision division with a
    /// handful of u64 × u64 multiplies. Falls back to Knuth's Algorithm D
    /// for general moduli.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    /// let P = U256::from_be_limbs([
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    /// ]);
    ///
    /// let a = U256::from_be_limbs([0, 0, 0, 7]);
    /// let b = U256::from_be_limbs([0, 0, 0, 6]);
    /// assert_eq!(a.mul_mod(&b, &P), U256::from_be_limbs([0, 0, 0, 42]));
    /// ```
    #[inline]
    pub fn mul_mod(&self, other: &U256, modulus: &U256) -> U256 {
        // Fast path: fused multiply-reduce for sparse primes.
        // Avoids U512 intermediate by inlining schoolbook + sparse_reduce.
        if let Some(comp) = modulus.has_u64_complement() {
            return self.mul_mod_sparse(other, modulus, comp);
        }

        // Generic path: widening multiply then Knuth Algorithm D
        let product = *self * *other;
        product.div_u256(modulus)
    }
}

#[cfg(test)]
mod ai_tests {
    use super::*;
    const P: U256 = U256::from_be_limbs([0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F]);

    /// Zero times anything is zero.
    #[test]
    fn zero_times_any() {
        let a = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        assert_eq!(U256::ZERO.mul_mod(&a, &P), U256::ZERO);
    }

    /// Anything times zero is zero.
    #[test]
    fn any_times_zero() {
        let a = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        assert_eq!(a.mul_mod(&U256::ZERO, &P), U256::ZERO);
    }

    /// Multiplicative identity: 1 * x = x.
    #[test]
    fn multiplicative_identity() {
        let a = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        assert_eq!(U256::ONE.mul_mod(&a, &P), a);
    }

    /// Multiplication is commutative: a * b = b * a.
    #[test]
    fn commutative() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        let b = U256::from_be_limbs([0, 0, 0, 6]);
        assert_eq!(a.mul_mod(&b, &P), b.mul_mod(&a, &P));
    }

    /// Small product: 7 * 6 = 42.
    #[test]
    fn small_product() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        let b = U256::from_be_limbs([0, 0, 0, 6]);
        assert_eq!(a.mul_mod(&b, &P), U256::from_be_limbs([0, 0, 0, 42]));
    }

    /// (P - 1)^2 mod P = 1 (since (P-1) = -1, and (-1)^2 = 1).
    #[test]
    fn p_minus_one_squared() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        assert_eq!(p_minus_1.mul_mod(&p_minus_1, &P), U256::ONE);
    }

    /// 2 * (P+1)/2 mod P = 1 (verifying known inverse of 2).
    #[test]
    fn two_times_inverse() {
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        let inv2 = U256::from_be_limbs([
            0x7FFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFF7FFFFE18,
        ]);
        assert_eq!(two.mul_mod(&inv2, &P), U256::ONE);
    }

    /// Unreduced input is reduced before multiplication.
    #[test]
    fn unreduced_input() {
        let max = U256::from_be_bytes([0xFF; 32]);
        assert_eq!(max.mul_mod(&U256::ONE, &P), max.modulo(&P));
    }

    /// Sparse path matches generic path for various inputs.
    #[test]
    fn sparse_matches_generic() {
        let cases = [
            (
                U256::from_be_limbs([0, 0, 0, 1]),
                U256::from_be_limbs([0, 0, 0, 1]),
            ),
            (
                U256::from_be_limbs([
                    0x123456789ABCDEF0,
                    0xFEDCBA9876543210,
                    0x1111222233334444,
                    0x5555666677778888,
                ]),
                U256::from_be_limbs([
                    0xAAAABBBBCCCCDDDD,
                    0xEEEEFFFF00001111,
                    0x2222333344445555,
                    0x6666777788889999,
                ]),
            ),
            (
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, 0xFFFFFFFEFFFFFC2E]),
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, 0xFFFFFFFEFFFFFC2E]),
            ),
        ];
        for (a, b) in &cases {
            let product = *a * *b;
            let generic = product.div_u256(&P);
            let sparse = product.sparse_reduce(&P, 0x1000003D1);
            assert_eq!(sparse, generic, "mismatch for {a:?} * {b:?}");
        }
    }

    /// Sparse reduce handles zero product correctly.
    #[test]
    fn sparse_reduce_zero() {
        let product = crate::u512::U512::from_be_limbs([0u64; 8]);
        assert_eq!(product.sparse_reduce(&P, 0x1000003D1), U256::ZERO);
    }

    /// Sparse reduce handles max product (stress test for overflow paths).
    #[test]
    fn sparse_reduce_max_product() {
        let product = crate::u512::U512::from_be_limbs([u64::MAX; 8]);
        let result = product.sparse_reduce(&P, 0x1000003D1);
        let generic = product.div_u256(&P);
        assert_eq!(result, generic);
    }
}
