//!
//! Provides degree-2 and degree-4 isogeny formulas on Montgomery curves
//! in projective (X:Z) coordinates, along with strategy-driven chain
//! evaluation for even-degree isogenies.

use super::{EcCurve, EcIsogEven, EcIsomorphism, EcPoint};
use crate::fp::{Fp2, FpBackend};

/// Kernel push-through constants for a degree-2 isogeny.
#[derive(Clone, Debug)]
pub struct EcKps2<L: crate::params::SecurityLevel> {
    pub k: EcPoint<L>,
}

/// Kernel push-through constants for a degree-4 isogeny.
#[derive(Clone, Debug)]
pub struct EcKps4<L: crate::params::SecurityLevel> {
    pub k: [EcPoint<L>; 3],
}

/// Compute a degree-2 isogeny with kernel generated by `P` (where `P != (0,0)`).
///
/// Returns `(kps, B)` where `B` is the codomain curve coefficient in
/// A24 form `(A+2C : 4C)` and `kps` contains the push-through constants
/// for evaluating the isogeny on other points.
#[inline]
pub fn xisog_2<L: FpBackend>(p: &EcPoint<L>) -> (EcKps2<L>, EcPoint<L>) {
    let bz = p.z.sqr();
    let bx = p.x.sqr();
    let bx = bz.sub(&bx);
    let kx = p.x.add(&p.z);
    let kz = p.x.sub(&p.z);
    (
        EcKps2 {
            k: EcPoint::new(kx, kz),
        },
        EcPoint::new(bx, bz),
    )
}

/// Compute a degree-2 singular isogeny (kernel through `(0:0)`).
///
/// Takes the current curve in A24 form. Returns `(kps, B24)` where `B24`
/// is the codomain in A24 form. Only used in the signing path.
#[inline]
pub fn xisog_2_singular<L: FpBackend>(a24: &EcPoint<L>) -> (EcKps2<L>, EcPoint<L>) {
    let four = Fp2::<L>::from_small(4);
    let t0 = a24.x.add(&a24.x);
    let t0 = t0.sub(&a24.z);
    let t0 = t0.add(&t0);
    let a24z_inv = a24.z.inv();
    let t0 = t0.mul(&a24z_inv);
    let kx = t0.clone();
    let b24x = t0.add(&t0);
    let t0 = t0.sqr();
    let t0 = t0.sub(&four);
    let t0 = t0.sqrt();
    let kz = t0.neg();
    let b24z = t0.add(&t0);
    let b24x = b24x.add(&b24z);
    let b24z = b24z.add(&b24z);
    (
        EcKps2 {
            k: EcPoint::new(kx, kz),
        },
        EcPoint::new(b24x, b24z),
    )
}

/// Compute a degree-4 isogeny with kernel generated by `P` such that
/// `[2]P != (0,0)`.
///
/// Returns `(kps, B)` where `B` is the codomain in A24 form and `kps`
/// contains constants for evaluating the isogeny via `xeval_4`.
#[inline]
pub fn xisog_4<L: FpBackend>(p: &EcPoint<L>) -> (EcKps4<L>, EcPoint<L>) {
    let k0x = p.x.sqr();
    let k0z = p.z.sqr();
    let k1x = k0z.add(&k0x);
    let k1z = k0z.sub(&k0x);
    let bx = k1x.mul(&k1z);
    let bz = k0z.sqr();

    let k2x = p.x.add(&p.z);
    let k1x_final = p.x.sub(&p.z);
    let k0x_final = k0z.add(&k0z);
    let k0x_final = k0x_final.add(&k0x_final);

    (
        EcKps4 {
            k: [
                EcPoint::new(k0x_final, k0z),
                EcPoint::new(k1x_final, k1z),
                EcPoint::new(k2x, Fp2::zero()),
            ],
        },
        EcPoint::new(bx, bz),
    )
}

/// Evaluate a degree-2 isogeny (non-singular) on a slice of points.
///
/// Writes the images into `out`. `out` and `q` may alias (same slice).
#[inline]
pub fn xeval_2<L: FpBackend>(out: &mut [EcPoint<L>], q: &[EcPoint<L>], kps: &EcKps2<L>) {
    for j in 0..q.len() {
        let t0 = q[j].x.add(&q[j].z);
        let t1 = q[j].x.sub(&q[j].z);
        let t2 = kps.k.x.mul(&t1);
        let t1 = kps.k.z.mul(&t0);
        let t0 = t2.add(&t1);
        let t1 = t2.sub(&t1);
        out[j].x = q[j].x.mul(&t0);
        out[j].z = q[j].z.mul(&t1);
    }
}

/// Evaluate a degree-2 isogeny on a slice of points, in place.
#[inline]
pub fn xeval_2_inplace<L: FpBackend>(points: &mut [EcPoint<L>], kps: &EcKps2<L>) {
    for pt in points.iter_mut() {
        let t0 = pt.x.add(&pt.z);
        let t1 = pt.x.sub(&pt.z);
        let t2 = kps.k.x.mul(&t1);
        let t1 = kps.k.z.mul(&t0);
        let t0 = t2.add(&t1);
        let t1 = t2.sub(&t1);
        pt.x = pt.x.mul(&t0);
        pt.z = pt.z.mul(&t1);
    }
}

/// Evaluate a degree-2 singular isogeny on a slice of points.
#[inline]
pub fn xeval_2_singular<L: FpBackend>(out: &mut [EcPoint<L>], q: &[EcPoint<L>], kps: &EcKps2<L>) {
    for i in 0..q.len() {
        let t0 = q[i].x.mul(&q[i].z);
        let t1 = kps.k.x.mul(&q[i].z);
        let t1 = t1.add(&q[i].x);
        let t1 = t1.mul(&q[i].x);
        out[i].x = q[i].z.sqr();
        out[i].x = out[i].x.add(&t1);
        out[i].z = t0.mul(&kps.k.z);
    }
}

/// Evaluate a degree-2 singular isogeny in place.
#[inline]
pub fn xeval_2_singular_inplace<L: FpBackend>(points: &mut [EcPoint<L>], kps: &EcKps2<L>) {
    for pt in points.iter_mut() {
        let t0 = pt.x.mul(&pt.z);
        let t1 = kps.k.x.mul(&pt.z);
        let t1 = t1.add(&pt.x);
        let t1 = t1.mul(&pt.x);
        let rx = pt.z.sqr();
        let rx = rx.add(&t1);
        let rz = t0.mul(&kps.k.z);
        pt.x = rx;
        pt.z = rz;
    }
}

/// Evaluate a degree-4 isogeny on a slice of points.
#[inline]
pub fn xeval_4<L: FpBackend>(out: &mut [EcPoint<L>], q: &[EcPoint<L>], kps: &EcKps4<L>) {
    for i in 0..q.len() {
        let t0 = q[i].x.add(&q[i].z);
        let t1 = q[i].x.sub(&q[i].z);
        let rx = t0.mul(&kps.k[1].x);
        let rz = t1.mul(&kps.k[2].x);
        let t0 = t0.mul(&t1);
        let t0 = t0.mul(&kps.k[0].x);
        let t1 = rx.add(&rz);
        let rz = rx.sub(&rz);
        let t1 = t1.sqr();
        let rz = rz.sqr();
        let rx = t0.add(&t1);
        let t0 = t0.sub(&rz);
        out[i].x = rx.mul(&t1);
        out[i].z = rz.mul(&t0);
    }
}

/// Evaluate a degree-4 isogeny in place.
#[inline]
pub fn xeval_4_inplace<L: FpBackend>(points: &mut [EcPoint<L>], kps: &EcKps4<L>) {
    for pt in points.iter_mut() {
        let t0 = pt.x.add(&pt.z);
        let t1 = pt.x.sub(&pt.z);
        let rx = t0.mul(&kps.k[1].x);
        let rz = t1.mul(&kps.k[2].x);
        let t0 = t0.mul(&t1);
        let t0 = t0.mul(&kps.k[0].x);
        let t1 = rx.add(&rz);
        let rz = rx.sub(&rz);
        let t1 = t1.sqr();
        let rz = rz.sqr();
        let rx = t0.add(&t1);
        let t0 = t0.sub(&rz);
        pt.x = rx.mul(&t1);
        pt.z = rz.mul(&t0);
    }
}

/// Strategy-driven evaluation of a 2ⁿ isogeny chain using degree-4 steps
/// with an optional trailing degree-2 step.
///
/// Modifies `curve` to the codomain and pushes all `points` through the
/// chain. Returns `None` on failure (degenerate kernel).
#[inline]
fn ec_eval_even_strategy<L: FpBackend>(
    curve: &mut EcCurve<L>,
    points: &mut [EcPoint<L>],
    kernel: &EcPoint<L>,
    isog_len: i32,
) -> Option<()> {
    use super::point::xdbl_a24;

    curve.normalize_a24();
    let mut a24 = curve.ac_to_a24();

    // Stack of remaining kernel points and associated remaining orders.
    // 32 entries suffices: ceil(log2(isog_len)) + 1 is at most ~20 for
    // any parameter set in the spec.
    let mut splits: [EcPoint<L>; 32] = core::array::from_fn(|_| EcPoint::identity());
    let mut todo: [u16; 32] = [0u16; 32];
    splits[0] = kernel.clone();
    todo[0] = isog_len as u16;

    let mut current: i32 = 0;

    // Chain of 4-isogenies
    for j in 0..(isog_len / 2) {
        while todo[current as usize] != 2 {
            current += 1;
            splits[current as usize] = splits[(current - 1) as usize].clone();
            let num_dbls =
                (todo[(current - 1) as usize] / 4 * 2 + todo[(current - 1) as usize] % 2) as i32;
            todo[current as usize] = todo[(current - 1) as usize] - num_dbls as u16;
            for _ in 0..num_dbls {
                splits[current as usize] = xdbl_a24(&splits[current as usize], &a24, false);
            }
        }

        if j == 0 {
            if !bool::from(a24.z.ct_is_one()) {
                return None;
            }
            if !bool::from(splits[current as usize].is_four_torsion(curve)) {
                return None;
            }
            let t = xdbl_a24(&splits[current as usize], &a24, false);
            if bool::from(t.x.ct_is_zero()) {
                return None;
            }
        }

        // Evaluate 4-isogeny
        let (kps4, new_a24) = xisog_4(&splits[current as usize]);
        a24 = new_a24;

        // Push splits[0..current] through the isogeny
        for i in 0..current as usize {
            xeval_4_inplace(&mut splits[i..i + 1], &kps4);
            todo[i] -= 2;
        }

        // Push user points through
        xeval_4_inplace(points, &kps4);

        current -= 1;
    }

    // Final 2-isogeny if isog_len is odd
    if isog_len % 2 != 0 {
        if isog_len == 1 && !bool::from(splits[0].is_two_torsion(curve)) {
            return None;
        }
        if bool::from(splits[0].x.ct_is_zero()) {
            return None;
        }

        let (kps2, new_a24) = xisog_2(&splits[0]);
        a24 = new_a24;

        xeval_2_inplace(points, &kps2);
    }

    // Convert back to (A:C) form
    *curve = EcCurve::from_a24(&a24);
    curve.is_a24_computed_and_normalized = false;
    Some(())
}

/// Evaluate an even-degree isogeny on a set of points.
///
/// Computes the codomain curve and pushes all `points` through the isogeny.
/// Returns `None` on failure (degenerate kernel).
#[inline]
pub fn ec_eval_even<L: FpBackend>(
    image: &mut EcCurve<L>,
    phi: &EcIsogEven<L>,
    points: &mut [EcPoint<L>],
) -> Option<()> {
    *image = phi.curve.clone();
    ec_eval_even_strategy(image, points, &phi.kernel, phi.length as i32)
}

/// Naive O(n^2) isogeny chain evaluation, one degree-2 step at a time.
///
/// Supports singular isogenies (kernel through `(0:0)`) when `special`
/// is true. Returns `None` on failure (kernel not of expected order).
#[inline]
pub fn ec_eval_small_chain<L: FpBackend>(
    curve: &mut EcCurve<L>,
    kernel: &EcPoint<L>,
    len: i32,
    points: &mut [EcPoint<L>],
    special: bool,
) -> Option<()> {
    let mut a24 = curve.ac_to_a24();
    let mut big_k = kernel.clone();

    for i in 0..len {
        let mut small_k = big_k.clone();
        for _ in 0..(len - i - 1) {
            small_k = super::point::xdbl_a24(&small_k, &a24, false);
        }

        if i == 0 && !bool::from(small_k.is_two_torsion(curve)) {
            return None;
        }

        if bool::from(small_k.x.ct_is_zero()) {
            if special {
                let (kps, b24) = xisog_2_singular(&a24);
                let mut big_k_arr = [big_k];
                xeval_2_singular_inplace(&mut big_k_arr, &kps);
                big_k = big_k_arr[0].clone();
                xeval_2_singular_inplace(points, &kps);
                a24 = b24;
            } else {
                return None;
            }
        } else {
            let (kps, new_a24) = xisog_2(&small_k);
            a24 = new_a24;

            let mut big_k_arr = [big_k];
            xeval_2_inplace(&mut big_k_arr, &kps);
            big_k = big_k_arr[0].clone();
            xeval_2_inplace(points, &kps);
        }
    }

    *curve = EcCurve::from_a24(&a24);
    curve.is_a24_computed_and_normalized = false;
    Some(())
}

/// Compute the isomorphism between two j-equivalent Montgomery curves.
///
/// Goes through Montgomery -> Short Weierstrass -> Short Weierstrass -> Montgomery.
/// The isomorphism maps `(X:Z)` on `from` to `(Nx*X + Nz*Z : D*Z)` on `to`.
/// Returns `None` if the isomorphism is degenerate (Nx or D is zero).
#[inline]
pub fn ec_isomorphism<L: FpBackend>(
    from: &EcCurve<L>,
    to: &EcCurve<L>,
) -> Option<EcIsomorphism<L>> {
    let t0 = from.a.mul(&from.c);
    let t1 = to.a.mul(&to.c);

    let t2 = t1.mul(&to.c); // toA*toC^2
    let t3 = t2.add(&t2);
    let t3 = t3.add(&t3);
    let t3 = t3.add(&t3);
    let t2 = t2.add(&t3); // 9*toA*toC^2
    let t3 = to.a.sqr();
    let t3 = t3.mul(&to.a); // toA^3
    let t3 = t3.add(&t3);
    let nx = t3.sub(&t2); // 2*toA^3 - 9*toA*toC^2

    let t2 = t0.mul(&from.a); // fromA^2*fromC
    let t3 = from.c.sqr();
    let t3 = t3.mul(&from.c); // fromC^3
    let t4 = t3.add(&t3);
    let t3 = t4.add(&t3); // 3*fromC^3
    let t3 = t3.sub(&t2); // 3*fromC^3 - fromA^2*fromC
    let nx = nx.mul(&t3); // lambda_x

    let t2 = t0.mul(&from.c); // fromA*fromC^2
    let t3 = t2.add(&t2);
    let t3 = t3.add(&t3);
    let t3 = t3.add(&t3);
    let t2 = t2.add(&t3); // 9*fromA*fromC^2
    let t3 = from.a.sqr();
    let t3 = t3.mul(&from.a); // fromA^3
    let t3 = t3.add(&t3);
    let d = t3.sub(&t2); // 2*fromA^3 - 9*fromA*fromC^2

    let t2 = t1.mul(&to.a); // toA^2*toC
    let t3 = to.c.sqr();
    let t3 = t3.mul(&to.c); // toC^3
    let t4 = t3.add(&t3);
    let t3 = t4.add(&t3); // 3*toC^3
    let t3 = t3.sub(&t2); // 3*toC^3 - toA^2*toC
    let d = d.mul(&t3); // lambda_z

    // Mont -> SW -> SW -> Mont
    let t0 = to.c.mul(&from.a);
    let t0 = t0.mul(&nx); // lambda_x * toC * fromA
    let t1 = from.c.mul(&to.a);
    let t1 = t1.mul(&d); // lambda_z * fromC * toA
    let nz = t0.sub(&t1); // lambda_x*toC*fromA - lambda_z*fromC*toA

    let t0 = from.c.mul(&to.c);
    let t1 = t0.add(&t0);
    let t0 = t0.add(&t1); // 3*fromC*toC
    let d = d.mul(&t0); // 3*lambda_z*fromC*toC
    let nx = nx.mul(&t0); // 3*lambda_x*fromC*toC

    if bool::from(nx.ct_is_zero() | d.ct_is_zero()) {
        return None;
    }
    Some(EcIsomorphism { nx, nz, d })
}

/// Apply an isomorphism to a point: `(X:Z) -> (Nx*X + Nz*Z : D*Z)`.
#[inline]
pub fn ec_iso_eval<L: FpBackend>(p: &mut EcPoint<L>, isom: &EcIsomorphism<L>) {
    let tmp = p.z.mul(&isom.nz);
    p.x = p.x.mul(&isom.nx);
    p.x = p.x.add(&tmp);
    p.z = p.z.mul(&isom.d);
}
