//! Modular square root for [`U256`] values.

use super::U256;

impl U256 {
    /// Computes the modular square root of `self` modulo `modulus`, returning
    /// `Some(r)` where `r² ≡ self (mod modulus)`, or `None` if `self` is not
    /// a quadratic residue.
    ///
    /// Only valid when `modulus ≡ 3 (mod 4)` (e.g. the secp256k1 field prime).
    /// Uses the direct formula `r = self^((modulus + 1) / 4) mod modulus`,
    /// which avoids the complexity of Tonelli-Shanks. The result is verified
    /// by squaring: if `r² ≢ self (mod modulus)`, `self` is not a quadratic
    /// residue and `None` is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    ///
    /// let p = U256::from_be_limbs([
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    /// ]);
    /// let four = U256::from_be_limbs([0, 0, 0, 4]);
    /// let root = four.sqrt_mod(&p).unwrap();
    /// assert_eq!(root.square_mod(&p), four);
    /// ```
    #[inline]
    pub fn sqrt_mod(&self, modulus: &U256) -> Option<U256> {
        // Exponent = (modulus + 1) / 4
        // Since modulus ≡ 3 (mod 4), modulus + 1 is divisible by 4.
        // Adding 1 to an odd number gives an even number, so no overflow
        // concern for 256-bit primes with leading bits set.
        let exp = (*modulus + U256::ONE) >> 2;

        let candidate = self.pow_mod(&exp, modulus);

        // Verify: candidate² ≡ self (mod modulus)
        if candidate.square_mod(modulus) == self.modulo(modulus) {
            Some(candidate)
        } else {
            None
        }
    }
}

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

    /// secp256k1 field prime.
    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFEFFFFFC2F,
    ]);

    /// secp256k1 generator y-coordinate.
    const GY: U256 = U256::from_be_limbs([
        0x483ADA7726A3C465,
        0x5DA4FBFC0E1108A8,
        0xFD17B448A6855419,
        0x9C47D08FFB10D4B8,
    ]);

    /// secp256k1 generator x-coordinate.
    const GX: U256 = U256::from_be_limbs([
        0x79BE667EF9DCBBAC,
        0x55A06295CE870B07,
        0x029BFCDB2DCE28D9,
        0x59F2815B16F81798,
    ]);

    /// Recovering y² from the generator x gives a root matching GY.
    #[test]
    fn generator_y_recovery() {
        // y² = x³ + 7 mod P
        let x2 = GX.square_mod(&P);
        let x3 = x2.mul_mod(&GX, &P);
        let seven = U256::from_be_limbs([0, 0, 0, 7]);
        let rhs = x3.add_mod(&seven, &P);

        let root = rhs.sqrt_mod(&P).expect("generator y² should be a QR");
        // sqrt_mod returns one of the two roots; GY or P - GY
        assert!(root == GY || root == P - GY);
        // Verify squaring back
        assert_eq!(root.square_mod(&P), rhs);
    }

    /// sqrt(4) mod P = 2.
    #[test]
    fn sqrt_of_four() {
        let four = U256::from_be_limbs([0, 0, 0, 4]);
        let root = four.sqrt_mod(&P).unwrap();
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        // Root is either 2 or P - 2
        assert!(root == two || root == P - two);
        assert_eq!(root.square_mod(&P), four);
    }

    /// sqrt(1) mod P = 1.
    #[test]
    fn sqrt_of_one() {
        let root = U256::ONE.sqrt_mod(&P).unwrap();
        assert!(root == U256::ONE || root == P - U256::ONE);
        assert_eq!(root.square_mod(&P), U256::ONE);
    }

    /// sqrt(0) mod P = 0.
    #[test]
    fn sqrt_of_zero() {
        let root = U256::ZERO.sqrt_mod(&P).unwrap();
        assert_eq!(root, U256::ZERO);
    }

    /// A known non-residue returns None.
    #[test]
    fn non_residue_returns_none() {
        // 3 is a quadratic non-residue mod P (Euler criterion: 3^((P-1)/2) ≡ -1)
        let three = U256::from_be_limbs([0, 0, 0, 3]);
        assert!(three.sqrt_mod(&P).is_none());
    }

    /// sqrt(9) mod P = 3.
    #[test]
    fn sqrt_of_nine() {
        let nine = U256::from_be_limbs([0, 0, 0, 9]);
        let root = nine.sqrt_mod(&P).unwrap();
        let three = U256::from_be_limbs([0, 0, 0, 3]);
        assert!(root == three || root == P - three);
        assert_eq!(root.square_mod(&P), nine);
    }

    /// Verifies that squaring the result always recovers the input.
    #[test]
    fn round_trip_small_residues() {
        for v in [4u64, 9, 16, 25, 36, 49, 64, 81, 100] {
            let val = U256::from_be_limbs([0, 0, 0, v]);
            let root = val.sqrt_mod(&P).expect(&format!("{v} should be a QR"));
            assert_eq!(root.square_mod(&P), val);
        }
    }
}
