#[cfg(debug_assertions)]
use crate::core::expressions::domain::Domain;
use crate::{
    core::{
        circuits::{
            arithmetic::sqrt,
            boolean::{boolean_value::BooleanValue, byte::Byte},
        },
        expressions::{
            curve_expr::CurveExpr,
            expr::Expr,
            field_expr::FieldExpr,
            other_expr::OtherExpr,
        },
        global_value::{global_expr_store::with_global_expr_store_as_local, value::FieldValue},
        ir_builder::{ExprStore, IRBuilder},
    },
    traits::{
        Equal,
        FromLeBits,
        FromLeBytes,
        GetBit,
        Invert,
        RandomBit,
        Reveal,
        Select,
        ToLeBytes,
        ToMontgomery,
    },
    utils::{
        curve_point::{Curve, CurvePoint},
        elliptic_curve::{
            AffineEdwardsPoint,
            ProjectiveEdwardsPoint,
            EDWARDS25519_D,
            INVSQRT_NEG_ONE_MINUS_D,
            SQRT_NEG_ONE,
        },
        field::{BaseField, ScalarField},
    },
};
use curve25519_dalek::RistrettoPoint;
use group::GroupEncoding;
use primitives::algebra::elliptic_curve::{Curve as AsyncMPCCurve, Curve25519Ristretto};
use std::ops::{Add, AddAssign, Mul, Neg, Sub};

/// A CurveValue is a secret-shared or public point in the cryptographic prime order group of
/// Curve25519. The backend for CurveValue is RistrettoPoint, which is an element of the
/// quotient group [2]E / E[4], and is internally represented by an EdwardsPoint.
/// Note: the order of the representative divides 4 * \ell, i.e., the EdwardsPoint
/// itself need not be of prime order.
#[derive(Debug, Clone, Copy)]
pub struct CurveValue(usize);

#[allow(non_snake_case)]
impl CurveValue {
    pub fn new(id: usize) -> Self {
        #[cfg(debug_assertions)]
        with_global_expr_store_as_local(|expr_store| {
            let _ = CurvePoint::unwrap(*expr_store.get_bounds(id));
        });
        Self(id)
    }
    pub fn from_expr(expr: Expr<usize>) -> Self {
        let id = with_global_expr_store_as_local(|expr_store| {
            let id = expr_store.new_expr(expr);
            #[cfg(debug_assertions)]
            let _ = CurvePoint::unwrap(*expr_store.get_bounds(id));
            id
        });
        Self(id)
    }

    pub fn get_id(&self) -> usize {
        self.0
    }

    pub fn expr(&self) -> Expr<usize> {
        with_global_expr_store_as_local(|expr_store| expr_store.get_expr(self.0).clone())
    }

    pub fn is_plaintext(&self) -> bool {
        with_global_expr_store_as_local(|expr_store| expr_store.get_is_plaintext(self.0))
    }

    pub fn multiscalar_mul(scalars: Vec<FieldValue<ScalarField>>, points: Vec<Self>) -> Self {
        assert_eq!(scalars.len(), points.len());
        scalars
            .into_iter()
            .zip(points)
            .map(|(scalar, point)| scalar * point)
            .reduce(|lhs, rhs| lhs + rhs)
            .unwrap()
    }

    /// Additive identity.
    pub fn identity() -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(
                expr_store,
                CurveExpr::Val(CurvePoint::identity()),
            )
        }))
    }

    /// Is self the additive identity.
    pub fn is_identity(&self) -> BooleanValue {
        let point = self.to_projective();
        // The PlaintextPointToExtendedEdwards gate (called by to_extended_public() in
        // to_projective()) asserts that self is not at infinity.
        // Hence the below is a sufficient criteria.
        point.Y.eq(point.Z)
    }

    /// Generator of the cryptographic prime order group.
    pub fn generator() -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(
                expr_store,
                CurveExpr::Val(CurvePoint::generator()),
            )
        }))
    }

    /// Outputs a doubly authenticated point, i.e., a random point in the prime order group that is
    /// both EC secret-shared and base field secret-shared (a ProjectiveEdwardsPoint).
    /// Da points are needed for the conversion between CurveValue and tuple/triple/quadruple of
    /// affine/projective/extended Edwards coordinates.
    pub fn da_point() -> (Self, ProjectiveEdwardsPoint<FieldValue<BaseField>>) {
        // The modulus of the scalar field is ell = 2^252 + eps, where eps =
        // 27742317777372353535851937790883648493 is a 125-bit number. One can show that
        // sampling a uniformly random 252-bit number has a distribution that is eps/ell ~
        // 2^-127 close to the uniform distribution on the scalar field.
        let rand_bits = (0..252)
            .map(|_| BooleanValue::random())
            .collect::<Vec<BooleanValue>>();
        let rand = FieldValue::<ScalarField>::from_le_bits(rand_bits.clone(), false);
        let da_point_curve = rand * Self::generator();
        let da_point_proj = ProjectiveEdwardsPoint::mul_bits_generator(rand_bits);
        (da_point_curve, da_point_proj)
    }

    /// Conversion from public CurveValue to public extended Edwards coordinates.
    /// Returns the coordinates of the prime order representative. Fails if self is at infinity.
    pub fn to_extended_public(
        self,
    ) -> (
        FieldValue<BaseField>,
        FieldValue<BaseField>,
        FieldValue<BaseField>,
        FieldValue<BaseField>,
    ) {
        assert!(self.is_plaintext());
        let (X_id, Y_id, Z_id, T_id) = with_global_expr_store_as_local(|expr_store| {
            (
                <IRBuilder as ExprStore<BaseField>>::push_other(
                    expr_store,
                    OtherExpr::PlaintextCurveToExtendedEdwards(self.0, 0),
                ),
                <IRBuilder as ExprStore<BaseField>>::push_other(
                    expr_store,
                    OtherExpr::PlaintextCurveToExtendedEdwards(self.0, 1),
                ),
                <IRBuilder as ExprStore<BaseField>>::push_other(
                    expr_store,
                    OtherExpr::PlaintextCurveToExtendedEdwards(self.0, 2),
                ),
                <IRBuilder as ExprStore<BaseField>>::push_other(
                    expr_store,
                    OtherExpr::PlaintextCurveToExtendedEdwards(self.0, 3),
                ),
            )
        });

        (
            FieldValue::<BaseField>::from_id(X_id),
            FieldValue::<BaseField>::from_id(Y_id),
            FieldValue::<BaseField>::from_id(Z_id),
            FieldValue::<BaseField>::from_id(T_id),
        )
    }

    /// Conversion from public extended Edwards coordinates to public CurveValue.
    /// Uses the isomorphism E / E[8] -> [2]E / E[4] induced by
    ///     E -> [2]E, P \mapsto [2^-1 mod \ell] \circ [2] P
    /// to convert points of order divisible by 8 to valid RistrettoPoints.
    /// If the coordinates do not define a point on the curve this function returns the identity.
    /// Note: invalid coordinates include the all-zero quadruple and points at infinity (the
    /// Edwards curve in extended coordinates does not have points at infinity).
    #[allow(dead_code)]
    fn from_extended_public(
        point: (
            FieldValue<BaseField>,
            FieldValue<BaseField>,
            FieldValue<BaseField>,
            FieldValue<BaseField>,
        ),
    ) -> Self {
        assert!(with_global_expr_store_as_local(|expr_store| {
            expr_store.get_is_plaintext(point.0.get_id())
                && expr_store.get_is_plaintext(point.1.get_id())
                && expr_store.get_is_plaintext(point.2.get_id())
                && expr_store.get_is_plaintext(point.3.get_id())
        }));
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<BaseField>>::push_other(
                expr_store,
                OtherExpr::CurveFromPlaintextExtendedEdwards(
                    point.0.get_id(),
                    point.1.get_id(),
                    point.2.get_id(),
                    point.3.get_id(),
                ),
            )
        }))
    }

    /// Conversion from CurveValue to projective coordinates.
    /// Note: this clears the 4-torsion component of the representative of self, i.e.,
    /// it returns a point of order ell.
    pub fn to_projective(self) -> ProjectiveEdwardsPoint<FieldValue<BaseField>> {
        let [X, Y, Z] = [0, 1, 2]
            .map(|t| FieldValue::from_expr(Expr::Other(OtherExpr::ToProjective(self.get_id(), t))));
        ProjectiveEdwardsPoint::new((X, Y, Z), true, true)
    }

    /// Conversion from projective coordinates to CurveValue.
    /// This clears the 8-torsion component of point, i.e., the double conversion
    /// ProjectiveEdwardsPoint -> CurveValue -> ProjectiveEdwardsPoint is the identity only for
    /// \ell-torsion points. If the coordinates do not define a non-singular point of the curve this
    /// function returns the identity.
    pub fn from_projective(point: ProjectiveEdwardsPoint<FieldValue<BaseField>>) -> Self {
        if with_global_expr_store_as_local(|expr_store| {
            expr_store.get_is_plaintext(point.X.get_id())
                && expr_store.get_is_plaintext(point.Y.get_id())
                && expr_store.get_is_plaintext(point.Z.get_id())
        }) {
            let (X, Y, Z) = point.inner();
            Self::from_extended_public((X * Z, Y * Z, Z * Z, X * Y))
        } else if point.is_on_curve {
            let (da_point_curve, da_point_proj) = Self::da_point();
            let masked = if point.is_ell_torsion {
                (point + da_point_proj).reveal()
            } else {
                // We must mask the 8-torsion component of point as well. Via the Chinese Remainder
                // Theorem, the below is equivalent to masking point with a random element of E(F_p)
                // with predefined \ell-torsion (da_point_proj).
                (point
                    // adding parantheses here, so that the full E(F_p) mask can be computed before point is known
                    + (da_point_proj + ProjectiveEdwardsPoint::random_eight_torsion_point::<BooleanValue>()
                    ))
                .reveal()
            };
            let (X, Y, Z) = masked.inner();
            let masked_curve = Self::from_extended_public((X * Z, Y * Z, Z * Z, X * Y));
            masked_curve - da_point_curve
        } else {
            let (X, Y, Z) = point.inner();
            let (X2, Y2, Z2) = (X * X, Y * Y, Z * Z);
            // Check if the coordinates define a non-singular point on the curve (the two points
            // at infinity (1:0:0) and (0:1:0) satisfy the curve equation but are singular points).
            let is_on_curve = (-X2 * Z2 + Y2 * Z2)
                .eq(Z2 * Z2 + FieldValue::from(BaseField::from_le_bytes(EDWARDS25519_D)) * X2 * Y2)
                & !Z.eq(FieldValue::<BaseField>::from(0));
            // If point is not on the curve we take the identity.
            let point = is_on_curve.select(point, ProjectiveEdwardsPoint::identity());
            let (da_point_curve, da_point_proj) = Self::da_point();
            // We must mask the 8-torsion component of point as well. Via the Chinese Remainder
            // Theorem, the below is equivalent to masking point with a random element of E(F_p)
            // with predefined \ell-torsion (da_point_proj).
            let masked = (point
                // adding parantheses here, so that the full E(F_p) mask can be computed before point is known
                + (da_point_proj + ProjectiveEdwardsPoint::random_eight_torsion_point::<BooleanValue>()))
            .reveal();
            let (X, Y, Z) = masked.inner();
            let masked_curve = Self::from_extended_public((X * Z, Y * Z, Z * Z, X * Y));
            masked_curve - da_point_curve
        }
    }

    /// Conversion from CurveValue to affine coordinates.
    /// Note: this clears the 4-torsion component of the representative of self, i.e.,
    /// it returns a point of order ell.
    pub fn to_affine(self) -> AffineEdwardsPoint<FieldValue<BaseField>> {
        self.to_projective().to_affine()
    }

    /// Conversion from affine coordinates to CurveValue.
    /// This clears the 8-torsion component of point, i.e., the double conversion
    /// AffineEdwardsPoint -> CurveValue -> AffineEdwardsPoint is the identity only for \ell-torsion
    /// points. If the coordinates do not define a point on the curve this function returns the
    /// identity.
    pub fn from_affine(point: AffineEdwardsPoint<FieldValue<BaseField>>) -> Self {
        Self::from_projective(point.to_projective())
    }

    pub fn compress(&self) -> CompressedCurveValue {
        if self.is_plaintext() {
            let compressed_ids = with_global_expr_store_as_local(|expr_store| {
                (0..256)
                    .map(|i| {
                        <IRBuilder as ExprStore<BaseField>>::push_other(
                            expr_store,
                            OtherExpr::CompressPlaintextPoint(self.0, i),
                        )
                    })
                    .collect::<Vec<usize>>()
            });
            let compressed_bits = compressed_ids
                .into_iter()
                .map(BooleanValue::new)
                .collect::<Vec<BooleanValue>>();

            CompressedCurveValue(
                compressed_bits
                    .chunks(8)
                    .map(|chunk| {
                        Byte::new(chunk.to_vec().try_into().unwrap_or_else(
                            |v: Vec<BooleanValue>| {
                                panic!("Expected a Vec of length 8 (found {})", v.len())
                            },
                        ))
                    })
                    .collect::<Vec<Byte<BooleanValue>>>()
                    .try_into()
                    .unwrap_or_else(|v: Vec<Byte<BooleanValue>>| {
                        panic!("Expected a Vec of length 32 (found {})", v.len())
                    }),
            )
        } else {
            let (mut X, mut Y, Z, T) = self.to_extended_public();

            let u1 = (Z + Y) * (Z - Y);
            let u2 = X * Y;
            // Ignore return value since this is always square
            let (_, invsqrt) = sqrt::<BaseField, BooleanValue, FieldValue<BaseField>>(
                (u1 * u2 * u2).invert(true),
                true,
            );
            // we need the non-negative square root
            let invsqrt = invsqrt.get_bit(0, false).select(-invsqrt, invsqrt);
            let i1 = invsqrt * u1;
            let i2 = invsqrt * u2;
            let z_inv = i1 * (i2 * T);
            let mut den_inv = i2;

            let sqrt_neg_one = FieldValue::from(BaseField::from_le_bytes(SQRT_NEG_ONE));
            let iX = X * sqrt_neg_one;
            let iY = Y * sqrt_neg_one;
            let ristretto_magic =
                FieldValue::from(BaseField::from_le_bytes(INVSQRT_NEG_ONE_MINUS_D));
            let enchanted_denominator = i1 * ristretto_magic;

            let rotate = (T * z_inv).get_bit(0, false);

            X = rotate.select(iY, X);
            Y = rotate.select(iX, Y);
            den_inv = rotate.select(enchanted_denominator, den_inv);

            Y = (X * z_inv).get_bit(0, false).select(-Y, Y);

            let mut s = den_inv * (Z - Y);
            s = s.get_bit(0, false).select(-s, s);

            CompressedCurveValue(s.to_le_bytes())
        }
    }
}

impl ToMontgomery for CurveValue {
    type Output = FieldValue<BaseField>;

    fn to_montgomery(
        self,
        is_expected_non_identity: bool,
    ) -> (FieldValue<BaseField>, FieldValue<BaseField>) {
        self.to_affine().to_montgomery(is_expected_non_identity)
    }
}

impl Curve for CurveValue {
    fn generator() -> Self {
        Self::generator()
    }
}

impl Default for CurveValue {
    fn default() -> Self {
        Self::identity()
    }
}

impl Add for CurveValue {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(
                expr_store,
                CurveExpr::Add(self.0, rhs.0),
            )
        }))
    }
}

impl AddAssign for CurveValue {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl Sub for CurveValue {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        self + (-rhs)
    }
}

impl Neg for CurveValue {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(expr_store, CurveExpr::Neg(self.0))
        }))
    }
}

impl Mul<CurveValue> for FieldValue<ScalarField> {
    type Output = CurveValue;

    fn mul(self, rhs: CurveValue) -> Self::Output {
        CurveValue(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(
                expr_store,
                CurveExpr::Mul(self.get_id(), rhs.0),
            )
        }))
    }
}

impl Mul<CurveValue> for ScalarField {
    type Output = CurveValue;

    fn mul(self, rhs: CurveValue) -> Self::Output {
        CurveValue(with_global_expr_store_as_local(|expr_store| {
            let id =
                <IRBuilder as ExprStore<ScalarField>>::push_field(expr_store, FieldExpr::Val(self));
            <IRBuilder as ExprStore<ScalarField>>::push_curve(expr_store, CurveExpr::Mul(id, rhs.0))
        }))
    }
}

impl Reveal for CurveValue {
    fn reveal(self) -> Self {
        // Note: when revealing a CurveValue, each individual secret-share gets serialized before
        // being transmitted, and serialization does compression of the RistrettoPoint. As such, a
        // transmitted secret-share does not leak the 4-torsion component of the original
        // secret-share. On the other hand, since node_i decompresses all received
        // secret-shares but does not compress and decompress its own share for the
        // reconstruction, it might hold a different (but equivalent, of course)
        // representative of the revealed point than node_j.
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(expr_store, CurveExpr::Reveal(self.0))
        }))
    }
}

impl From<CurvePoint> for CurveValue {
    fn from(value: CurvePoint) -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<ScalarField>>::push_curve(expr_store, CurveExpr::Val(value))
        }))
    }
}

impl From<RistrettoPoint> for CurveValue {
    fn from(value: RistrettoPoint) -> Self {
        Self::from(CurvePoint::new(
            <Curve25519Ristretto as AsyncMPCCurve>::Point::from_bytes(&value.to_bytes()).unwrap(),
        ))
    }
}

#[derive(Debug, Clone, Copy)]
pub struct CompressedCurveValue([Byte<BooleanValue>; 32]);

impl CompressedCurveValue {
    pub fn to_bytes(self) -> [Byte<BooleanValue>; 32] {
        self.0
    }
}
