use crate::{
    core::{
        bounds::FieldBounds,
        circuits::{
            arithmetic::sqrt,
            boolean::{
                boolean_value::{Boolean, BooleanValue},
                byte::Byte,
            },
            general::edwards_mul::ProjectiveEdwards25519MultiplicationCircuit,
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{Equal, FromLeBytes, GetBit, GreaterEqual, Select},
    utils::{
        elliptic_curve::{
            AffineEdwardsPoint,
            EDWARDS25519_D,
            EDWARDS25519_D_TWIST,
            EIGHT_INV_MOD_ELL,
            F25519,
            SQRT_NEG_ONE,
        },
        field::BaseField,
        used_field::UsedField,
    },
};
use ff::Field;
use std::ops::Shr;

/// Encodes a base field element into a tuple of points on Edwards25519 x Edwards25519_twist (in
/// the large prime-order subgroups). To make the decoding work we require the base field element to
/// be at most 2^128.
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub struct BaseFieldToEdwards25519WithTwistEncodingCircuit {
    d: BaseField,
    d_twist: BaseField,
}

impl Default for BaseFieldToEdwards25519WithTwistEncodingCircuit {
    fn default() -> Self {
        Self::new()
    }
}

impl BaseFieldToEdwards25519WithTwistEncodingCircuit {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            d: BaseField::from_le_bytes(EDWARDS25519_D),
            d_twist: BaseField::from_le_bytes(EDWARDS25519_D_TWIST),
        }
    }

    #[cfg(test)]
    pub fn get_d(&self) -> BaseField {
        self.d
    }

    #[cfg(test)]
    pub fn get_d_twist(&self) -> BaseField {
        self.d_twist
    }
}

#[allow(non_snake_case, dead_code)]
impl BaseFieldToEdwards25519WithTwistEncodingCircuit {
    /// Sends the x-coordinate of a point to the square of the y-coordinate of the point
    /// (provided the point lies on the twisted Edwards curve).
    fn x_to_y2<T: F25519>(&self, x: T) -> T {
        let x2 = x * x;
        // here, the base is non-zero (because d * x2 is a quadratic non-residue)
        (T::from(1) + x2) * (T::from(1) - T::from(self.d) * x2).invert(true)
    }

    /// Encodes the message m into a tuple (Edwards25519, Edwards25519_twist).
    fn encode<
        B: Boolean + Select<T, T, T>,
        T: F25519 + GetBit<Output = B> + Equal<T, Output = B> + From<B>,
    >(
        &self,
        m: T,
    ) -> (AffineEdwardsPoint<T>, (T, T)) {
        // TODO: can we add the requirement on the input being at most 2^128 here?

        // Unlike BaseFieldToEdwards25519EncodingCircuit, where we encode the message into the
        // y-coordinate of a point on Edwards25519, here we must encode it into the
        // x-coordinate. The reason is that for any base field element x, either x is on the curve
        // or x^-1 is on the twist of the curve. For the y-coordinate, a similar relation would be:
        // y is on the curve if and only if sqrt(-d)*y is on the twist of the curve. However, the
        // latter lives in a degree-2 extension and we want to avoid computing it.

        // we need to add an offset (as m = 0 and m = 1 lead to 2- and 4-torsion points
        // respectively)
        let x = m + T::from(2);
        // all the below functions that multiplicatively mask their argument (such as g, sqrt,
        // safe_inverse) can omit the is_zero circuit.
        let y2 = Self::x_to_y2(self, x);
        let (is_on_curve, y) = sqrt::<BaseField, _, _>(y2, true);
        let x_twist = x.invert(true);
        // Note: x_to_y2(x_twist, d_twist) = -d * x_to_y2(x, d), since the F_{p^2}-isomorphism
        // between the curve and its twist sends (x, y) to (x^-1, sqrt(-d)*y).
        let (_, y_twist) = sqrt::<BaseField, _, _>(-T::from(self.d) * y2, true);

        let P = AffineEdwardsPoint::new(
            (
                is_on_curve.select(x, T::from(0)),
                is_on_curve.select(y, T::from(1)),
            ),
            true,
            false,
        );
        let P_twist = (
            is_on_curve.select(T::from(0), x_twist),
            is_on_curve.select(T::from(1), y_twist),
        );
        let P_twist_proj = (P_twist.0, P_twist.1, T::from(1));

        // TODO: here, we will switch from base field secret-sharing to EC point secret-sharing.
        // When generating random curve points Lambda and Lambda_twist in the prime-order subgroups
        // (both in EC point and base field secret-sharing) one must not forget to add a 8-torsion
        // resp. 4-torsion part to the base field secret-shared mask.

        let mul_circuit = ProjectiveEdwards25519MultiplicationCircuit::new();
        let res_proj = mul_circuit.mul_twist_str("001", P_twist_proj);
        let Z_inv = res_proj.2.invert(true);
        (
            P.to_projective().mul_str("0001").to_affine(),
            (res_proj.0 * Z_inv, res_proj.1 * Z_inv),
        )
    }
}

impl ArithmeticCircuit<BaseField> for BaseFieldToEdwards25519WithTwistEncodingCircuit {
    #[allow(non_snake_case)]
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        if x.len() != 1 {
            panic!("Base field to Edwards25519 with twist encoding requires input of length 1");
        }
        if x[0] > BaseField::power_of_two(128) {
            EvalFailure::err_ub("Input must be at most 2^128")
        } else {
            let (P, P_twist) = Self::encode::<bool, _>(self, x[0]);
            Ok(vec![P.x, P.y, P_twist.0, P_twist.1])
        }
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![
            FieldBounds::All,
            FieldBounds::All,
            FieldBounds::All,
            FieldBounds::All,
        ]
    }

    #[allow(non_snake_case)]
    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        if vals.len() != 1 {
            panic!("Base field to Edwards25519 with twist encoding requires input of length 1");
        }
        // TODO: add some diagnostics
        // let bounds = BaseField::bounds_to_field_bounds(vals[0].bounds());
        // if bounds.has_negatives() || bounds.unsigned_max() > BaseField::power_of_two(128) {
        //     EvalFailure::ub("Input must be at most 2^128")
        // }
        let (P, P_twist) = Self::encode::<BooleanValue, _>(self, vals[0]);
        vec![P.x, P.y, P_twist.0, P_twist.1]
    }
}

/// Decodes a base field element which was encoded via
/// BaseFieldToEdwards25519WithTwistEncodingCircuit.
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub struct Edwards25519WithTwistToBaseFieldDecodingCircuit {
    d: BaseField,
    d_twist: BaseField,
}

impl Default for Edwards25519WithTwistToBaseFieldDecodingCircuit {
    fn default() -> Self {
        Self::new()
    }
}

impl Edwards25519WithTwistToBaseFieldDecodingCircuit {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            d: BaseField::from_le_bytes(EDWARDS25519_D),
            d_twist: BaseField::from_le_bytes(EDWARDS25519_D_TWIST),
        }
    }
}

#[allow(non_snake_case, dead_code)]
impl Edwards25519WithTwistToBaseFieldDecodingCircuit {
    /// Computes the maximum between the minimum x-coordinate of the 8-torsion coset of a point P
    /// and the minimum x-coordinate of the 4-torsion coset of a point P_twist.
    /// Note that Q is any of the four rational points on Edwards25519 of order 8.
    fn recover_message<
        B: Boolean + Select<T, T, T>,
        T: F25519 + Shr<usize, Output = T> + GreaterEqual<T, Output = B>,
    >(
        P: AffineEdwardsPoint<T>,
        P_plus_Q: AffineEdwardsPoint<T>,
        P_twist: (T, T),
    ) -> T {
        let sqrt_neg_one = T::from(BaseField::from_le_bytes(SQRT_NEG_ONE));
        // Note: the action of the rational 4-torsion on a point (x, y) yields the coset {(x, y),
        // (sqrt(-1)*y, sqrt(-1)*x), (-x, -y), (-sqrt(-1)*y, -sqrt(-1)*x)}. That is, the set of
        // x-coordinates of that coset is precisely {x, -x, sqrt(-1)*y, -sqrt(-1)*y}.
        let mut x_coordinates = vec![
            P.x,
            -P.x,
            sqrt_neg_one * P.y,
            -sqrt_neg_one * P.y,
            P_plus_Q.x,
            -P_plus_Q.x,
            sqrt_neg_one * P_plus_Q.y,
            -sqrt_neg_one * P_plus_Q.y,
        ];
        let mut x_coordinates_twist = vec![
            P_twist.0,
            -P_twist.0,
            sqrt_neg_one * P_twist.1,
            -sqrt_neg_one * P_twist.1,
        ];

        // Note: P_twist is either the identity (in this case x_coordinates_twist equals [0, 0,
        // sqrt(-1), -sqrt(-1)]) or it is a non-zero element of large prime order (in this
        // case x_coordinates_twist equals [x_twist, -x_twist, sqrt(-1)*y_twist,
        // -sqrt(-1)*y_twist], and the first 4 elements of x_coordinates are [0, 0, sqrt(-1),
        // -sqrt(-1)]).

        // Since the encoding on the twist involved inverting finite field elements, we have to
        // invert the elements of x_coordinates_twist before we recover the message via the minimum
        // circuit. However, in about 50% of the cases, this will imply inverting the 0
        // element, which reveals that the message was encoded into a point on the curve and
        // not on the twist. With the below we obfuscate this information, without the need
        // of calling the is_zero circuit.

        // If the encoding happened into the twist then this doesn't change x_coordinates_twist
        // (we're adding 0's). However, if the encoding happened on the original curve then
        // this removes the 0's in x_coordinates_twist, which we obliviously add again via
        // could_be_one^2 - 1.
        x_coordinates_twist[0] = x_coordinates_twist[0] + x_coordinates[0];
        x_coordinates_twist[1] = x_coordinates_twist[1] + x_coordinates[1];
        let could_be_one = x_coordinates_twist[2] * sqrt_neg_one.invert(true);
        let mut x_coordinates_twist_inv = x_coordinates_twist
            .into_iter()
            .map(|c| {
                //all elements that we invert are non-zero
                c.invert(true)
            })
            .collect::<Vec<T>>();
        // TODO: for the below one could also append could_be_one - 1, but I have the feeling
        // that could_be_one^2 - 1 has less chance of being a 'small' element.
        x_coordinates_twist_inv.push(could_be_one * could_be_one - T::from(1));

        let min_x = x_coordinates
            .split_off(1)
            .into_iter()
            .fold(x_coordinates[0], |current, cand| {
                (cand.lt(current)).select(cand, current)
            });
        let min_x_twist = x_coordinates_twist_inv
            .split_off(1)
            .into_iter()
            .fold(x_coordinates_twist_inv[0], |current, cand| {
                (cand.lt(current)).select(cand, current)
            });

        // one of the two minima is 0, and we subtract the offset
        min_x + min_x_twist - T::from(2)
    }

    /// Decodes a base field element which was encoded via
    /// BaseFieldToEdwards25519WithTwistEncodingCircuit.
    fn decode<
        B: Boolean + Select<T, T, T>,
        T: F25519 + Shr<usize, Output = T> + GreaterEqual<T, Output = B>,
    >(
        &self,
        P: AffineEdwardsPoint<T>,
        P_twist: (T, T),
    ) -> T {
        let P_prime = P
            .to_projective()
            .mul_bits(
                EIGHT_INV_MOD_ELL
                    .into_iter()
                    .flat_map(|byte| Byte::<bool>::from(byte).to_vec())
                    .collect::<Vec<bool>>(),
            )
            .to_affine();
        let P_prime_plus_Q = P_prime + AffineEdwardsPoint::eight_torsion_point();

        let P_twist_proj = (P_twist.0, P_twist.1, T::from(1));
        // lsb-to-msb binary expansion of 4^{-1} modulo the order of the prime order subgroup
        let four_inv_bin = "0110100001000010111100010010111000011010110101100010011111011110011111010010100100110001110100000100110010010001100011010000011111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111101";
        let mul_circuit = ProjectiveEdwards25519MultiplicationCircuit::new();
        let P_twist_prime = mul_circuit.mul_twist_str(four_inv_bin, P_twist_proj);
        let Z_inv = P_twist_prime.2.invert(true);

        Self::recover_message(
            P_prime,
            P_prime_plus_Q,
            (P_twist_prime.0 * Z_inv, P_twist_prime.1 * Z_inv),
        )
    }
}

impl ArithmeticCircuit<BaseField> for Edwards25519WithTwistToBaseFieldDecodingCircuit {
    #[allow(non_snake_case)]
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        if x.len() != 4 {
            panic!("Edwards25519 with twist to Base field decoding requires input of length 4");
        }
        let P = AffineEdwardsPoint::new((x[0], x[1]), false, false);
        let P_twist = (x[2], x[3]);
        let m = Self::decode::<bool, _>(self, P, P_twist);
        if m > BaseField::power_of_two(128) {
            EvalFailure::err_ub("Circuit only supported for decoded messages of at most 2^128")
        } else {
            Ok(vec![m])
        }
    }

    fn bounds(&self, _bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        vec![FieldBounds::new(
            BaseField::ZERO,
            BaseField::power_of_two(128),
        )]
    }

    #[allow(non_snake_case)]
    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        if vals.len() != 4 {
            panic!("Edwards25519 with twist to Base field decoding requires input of length 4");
        }
        let P = AffineEdwardsPoint::new((vals[0], vals[1]), false, false);
        let P_twist = (vals[2], vals[3]);
        vec![Self::decode::<BooleanValue, _>(self, P, P_twist)]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit;
    use rand::Rng;

    #[test]
    #[allow(non_snake_case)]
    fn test_encoding_decoding() {
        let rng = &mut crate::utils::test_rng::get();
        let encoding_circuit = BaseFieldToEdwards25519WithTwistEncodingCircuit::new();
        let decoding_circuit = Edwards25519WithTwistToBaseFieldDecodingCircuit::new();
        for _ in 0..4 {
            for magnitude in [8, 16, 32, 64, 128] {
                let m = BaseField::gen_inclusive_range(
                    &mut *rng,
                    BaseField::ZERO,
                    BaseField::power_of_two(magnitude),
                );
                let (P, P_twist) = encoding_circuit.encode::<bool, _>(m);
                let (x, y) = P.inner();
                let x2 = x * x;
                let y2 = y * y;
                assert_eq!(
                    -x2 + y2,
                    BaseField::ONE + encoding_circuit.get_d() * x2 * y2
                );
                let x_twist2 = P_twist.0 * P_twist.0;
                let y_twist2 = P_twist.1 * P_twist.1;
                assert_eq!(
                    -x_twist2 + y_twist2,
                    BaseField::ONE + encoding_circuit.get_d_twist() * x_twist2 * y_twist2
                );
                let m_dec = decoding_circuit.decode::<bool, _>(P, P_twist);
                assert_eq!(m_dec, m);
            }
        }
    }

    impl TestedArithmeticCircuit<BaseField> for BaseFieldToEdwards25519WithTwistEncodingCircuit {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            Self::new()
        }

        fn gen_n_inputs<R: Rng + ?Sized>(&self, _rng: &mut R) -> usize {
            1
        }

        fn extra_checks(&self, inputs: Vec<BaseField>, outputs: Vec<BaseField>) {
            let decoding_circuit = Edwards25519WithTwistToBaseFieldDecodingCircuit::new();
            let decoded = decoding_circuit.eval(outputs).unwrap();
            assert_eq!(inputs, decoded);
        }
    }

    #[test]
    fn tested_base_field_to_edwards25519_with_twist_encoding() {
        BaseFieldToEdwards25519WithTwistEncodingCircuit::test(1, 4)
    }
}
