//! Inverse extraction from completed [`Divstep`] state.
use super::Divstep;
use crate::u256::U256;

impl Divstep {
    /// Extracts the modular inverse from the completed divstep state.
    ///
    /// After [`run()`](Divstep::run), `g = 0` and `|f| = 1`. The invariant
    /// `d * input ≡ f (mod modulus)` holds throughout, so when `f = ±1`:
    ///
    /// - `f = +1`: the inverse is `d`
    /// - `f = -1`: the inverse is `modulus - d` (negation)
    ///
    /// Returns `None` if `g ≠ 0`, `|f| ≠ 1`, or the result is zero
    /// (input was zero or a multiple of modulus).
    pub(crate) fn inverse(self) -> Option<U256> {
        if !self.g.is_zero() {
            return None;
        }

        if self.f != U256::ONE {
            return None;
        }

        // d * input ≡ f (mod modulus).
        // If f = +1 (f_neg=false), inverse = d.
        // If f = -1 (f_neg=true), inverse = -d = modulus - d.
        let result = if self.f_neg {
            U256::ZERO.sub_mod(&self.d, &self.modulus)
        } else {
            self.d
        };

        if result.is_zero() {
            None
        } else {
            Some(result)
        }
    }
}

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

    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    ]);

    /// Inverse of 7 matches Fermat oracle.
    #[test]
    fn inverse_of_seven() {
        let inv = Divstep::new(P, U256::from_be_limbs([0, 0, 0, 7]))
            .run()
            .inverse()
            .unwrap();
        let fermat = U256::from_be_limbs([0, 0, 0, 7]).mod_inv_fermat(&P).unwrap();
        assert_eq!(inv, fermat);
    }

    /// Inverse of 1 is 1.
    #[test]
    fn inverse_of_one() {
        let inv = Divstep::new(P, U256::ONE)
            .run()
            .inverse()
            .unwrap();
        assert_eq!(inv, U256::ONE);
    }

    /// Zero has no inverse.
    #[test]
    fn zero_has_no_inverse() {
        let result = Divstep::new(P, U256::ZERO)
            .run()
            .inverse();
        assert!(result.is_none());
    }

    /// Roundtrip: a * inv(a) mod P = 1.
    #[test]
    fn roundtrip() {
        let a = U256::from_be_limbs([0, 0, 0, 42]);
        let inv = Divstep::new(P, a).run().inverse().unwrap();
        assert_eq!(a.mul_mod(&inv, &P), U256::ONE);
    }

    /// Generator x-coordinate matches Lehmer inverse.
    #[test]
    fn generator_matches_lehmer() {
        let gx = U256::from_be_limbs([
            0x79BE667EF9DCBBAC, 0x55A06295CE870B07,
            0x029BFCDB2DCE28D9, 0x59F2815B16F81798,
        ]);
        let divstep_inv = Divstep::new(P, gx).run().inverse().unwrap();
        let lehmer_inv = gx.mod_inv(&P).unwrap();
        assert_eq!(divstep_inv, lehmer_inv);
    }

    /// Inverse of P-1 is P-1.
    #[test]
    fn inverse_of_p_minus_one() {
        let p_minus_1 = P - U256::ONE;
        let inv = Divstep::new(P, p_minus_1).run().inverse().unwrap();
        assert_eq!(inv, p_minus_1);
    }

    /// Inverse of 2 verified by roundtrip.
    #[test]
    fn inverse_of_two() {
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        let inv = Divstep::new(P, two).run().inverse().unwrap();
        assert_eq!(two.mul_mod(&inv, &P), U256::ONE);
    }

    /// Multiple values match Fermat oracle.
    #[test]
    fn batch_matches_fermat() {
        let values = [
            U256::from_be_limbs([0, 0, 0, 3]),
            U256::from_be_limbs([0, 0, 0, 100]),
            U256::from_be_limbs([0, 0, 0, 65537]),
            U256::from_be_limbs([
                0x79BE667EF9DCBBAC, 0x55A06295CE870B07,
                0x029BFCDB2DCE28D9, 0x59F2815B16F81798,
            ]),
        ];
        for v in &values {
            let divstep = Divstep::new(P, *v).run().inverse().unwrap();
            let fermat = v.mod_inv_fermat(&P).unwrap();
            assert_eq!(divstep, fermat, "mismatch for {v:?}");
        }
    }
}
