//! Variable-time fused multiply-and-reduce for sparse primes.
use super::U256;

impl U256 {
    /// Computes `(self * other) mod modulus` for sparse primes in a single
    /// fused pass: schoolbook multiply directly into sparse reduce.
    ///
    /// Instead of materializing a [`U512`] and then reducing, this function
    /// keeps the 8-limb product in local variables and immediately applies
    /// the folding reduction `lo + hi * comp`. This eliminates the U512
    /// struct overhead and gives LLVM a single large basic block for
    /// optimal register allocation and instruction scheduling.
    ///
    /// `comp` is the sparse prime complement: `2^256 - modulus`, which must
    /// fit in a single `u64`.
    ///
    /// # 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_sparse_vt(&b, &P, 0x1000003D1), U256::from_be_limbs([0, 0, 0, 42]));
    /// ```
    #[inline]
    pub fn mul_mod_sparse_vt(&self, other: &U256, modulus: &U256, comp: u64) -> U256 {
        let a = self.0;
        let b = other.0;

        // === Phase 1: Schoolbook 4×4 multiply (fully unrolled) ===

        // Row 0: a[0] * b[0..3]
        let p = (a[0] as u128) * (b[0] as u128);
        let r0 = p as u64;
        let carry = p >> 64;

        let p = (a[0] as u128) * (b[1] as u128) + carry;
        let r1 = p as u64;
        let carry = p >> 64;

        let p = (a[0] as u128) * (b[2] as u128) + carry;
        let r2 = p as u64;
        let carry = p >> 64;

        let p = (a[0] as u128) * (b[3] as u128) + carry;
        let r3 = p as u64;
        let r4 = (p >> 64) as u64;

        // Row 1: a[1] * b[0..3], accumulate into r[1..5]
        let p = (a[1] as u128) * (b[0] as u128) + (r1 as u128);
        let r1 = p as u64;
        let carry = p >> 64;

        let p = (a[1] as u128) * (b[1] as u128) + (r2 as u128) + carry;
        let r2 = p as u64;
        let carry = p >> 64;

        let p = (a[1] as u128) * (b[2] as u128) + (r3 as u128) + carry;
        let r3 = p as u64;
        let carry = p >> 64;

        let p = (a[1] as u128) * (b[3] as u128) + (r4 as u128) + carry;
        let r4 = p as u64;
        let r5 = (p >> 64) as u64;

        // Row 2: a[2] * b[0..3], accumulate into r[2..6]
        let p = (a[2] as u128) * (b[0] as u128) + (r2 as u128);
        let r2 = p as u64;
        let carry = p >> 64;

        let p = (a[2] as u128) * (b[1] as u128) + (r3 as u128) + carry;
        let r3 = p as u64;
        let carry = p >> 64;

        let p = (a[2] as u128) * (b[2] as u128) + (r4 as u128) + carry;
        let r4 = p as u64;
        let carry = p >> 64;

        let p = (a[2] as u128) * (b[3] as u128) + (r5 as u128) + carry;
        let r5 = p as u64;
        let r6 = (p >> 64) as u64;

        // Row 3: a[3] * b[0..3], accumulate into r[3..7]
        let p = (a[3] as u128) * (b[0] as u128) + (r3 as u128);
        let r3 = p as u64;
        let carry = p >> 64;

        let p = (a[3] as u128) * (b[1] as u128) + (r4 as u128) + carry;
        let r4 = p as u64;
        let carry = p >> 64;

        let p = (a[3] as u128) * (b[2] as u128) + (r5 as u128) + carry;
        let r5 = p as u64;
        let carry = p >> 64;

        let p = (a[3] as u128) * (b[3] as u128) + (r6 as u128) + carry;
        let r6 = p as u64;
        let r7 = (p >> 64) as u64;

        // === Phase 2: Sparse reduce (inline) ===
        // lo = [r0, r1, r2, r3], hi = [r4, r5, r6, r7]
        // result ≡ lo + hi * comp (mod modulus)

        let c = comp as u128;

        // Round 1: fused multiply-accumulate of hi * comp + lo
        let sum = (r0 as u128) + c * (r4 as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (r1 as u128) + c * (r5 as u128) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (r2 as u128) + c * (r6 as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (r3 as u128) + c * (r7 as u128) + carry;
        let r3 = sum as u64;
        let overflow = (sum >> 64) as u64;

        // Round 2: fold remaining overflow (typically ~33 bits for secp256k1)
        if overflow == 0 {
            return U256([r0, r1, r2, r3]).modulo(modulus);
        }

        let correction = (overflow as u128) * c;

        let sum = (r0 as u128) + (correction as u64 as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (r1 as u128) + (correction >> 64) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (r2 as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (r3 as u128) + carry;
        let r3 = sum as u64;
        let overflow2 = (sum >> 64) as u64;

        let mut result = U256([r0, r1, r2, r3]);

        if (overflow2 != 0) | (result >= *modulus) {
            result = result.overflowing_sub(modulus).0;
        }

        result.modulo(modulus)
    }
}

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

    /// 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_sparse_vt(&b, &P, COMP), U256::from_be_limbs([0, 0, 0, 42]));
    }

    /// 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_sparse_vt(&a, &P, COMP), U256::ZERO);
    }

    /// (P - 1)^2 mod P = 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_sparse_vt(&p_minus_1, &P, COMP), U256::ONE);
    }

    /// Matches mul_mod for diverse inputs.
    #[test]
    fn matches_mul_mod() {
        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]),
            ),
            (
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
            ),
        ];
        for (a, b) in &cases {
            let expected = a.mul_mod(b, &P);
            let fused = a.mul_mod_sparse_vt(b, &P, COMP);
            assert_eq!(fused, expected, "mismatch for {a:?} * {b:?}");
        }
    }

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

    /// Commutativity.
    #[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_sparse_vt(&b, &P, COMP),
            b.mul_mod_sparse_vt(&a, &P, COMP)
        );
    }
}
