//! Affine curve points.

#![allow(clippy::op_ref)]

use crate::{PrimeCurveParams, ProjectivePoint};
use core::borrow::Borrow;
use elliptic_curve::{
    Error, FieldBytes, Generate, PublicKey, Result, Scalar,
    bigint::modular::Retrieve,
    ctutils::{self, CtGt as _, CtSelect as _},
    ff::{Field, PrimeField},
    group::{CurveAffine, GroupEncoding},
    ops::{Double, Mul, MulVartime, Neg},
    point::{AffineCoordinates, DecompactPoint, DecompressPoint, NonIdentity},
    rand_core::{TryCryptoRng, TryRng},
    sec1::{self, CompressedPoint, FromSec1Point, Sec1Point, ToCompactSec1Point, ToSec1Point},
    subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption},
    zeroize::DefaultIsZeroes,
};

#[cfg(feature = "serde")]
use serdect::serde::{Deserialize, Serialize, de, ser};

/// Point on a Weierstrass curve in affine coordinates.
#[derive(Clone, Copy, Debug)]
pub struct AffinePoint<C: PrimeCurveParams> {
    /// x-coordinate
    pub(crate) x: C::FieldElement,

    /// y-coordinate
    pub(crate) y: C::FieldElement,

    /// Is this point the point at infinity? 0 = no, 1 = yes
    ///
    /// This is a proxy for [`Choice`], but uses `u8` instead to permit `const`
    /// constructors for `IDENTITY` and `GENERATOR`.
    pub(crate) infinity: u8,
}

impl<C> AffinePoint<C>
where
    C: PrimeCurveParams,
{
    /// Additive identity of the group a.k.a. the point at infinity.
    pub const IDENTITY: Self = Self {
        x: C::FieldElement::ZERO,
        y: C::FieldElement::ZERO,
        infinity: 1,
    };

    /// Base point of the curve.
    pub const GENERATOR: Self = Self {
        x: C::GENERATOR.0,
        y: C::GENERATOR.1,
        infinity: 0,
    };

    /// Is this point the point at infinity?
    pub fn is_identity(&self) -> Choice {
        Choice::from(self.infinity)
    }

    /// Conditionally negate [`AffinePoint`] for use with point compaction.
    fn to_compact(self) -> Self {
        let neg_self = -self;
        let gt = self.y.retrieve().ct_gt(&neg_self.y.retrieve());
        let y = C::FieldElement::conditional_select(&self.y, &neg_self.y, gt.into());
        Self {
            x: self.x,
            y,
            infinity: self.infinity,
        }
    }

    /// Internal RNG that avoids a `TryCryptoRng` bound so we can use it with `group`.
    ///
    /// TODO(tarcieri): find some way to avoid this?
    pub(crate) fn try_random<R: TryRng + ?Sized>(
        rng: &mut R,
    ) -> core::result::Result<Self, R::Error> {
        let mut bytes = FieldBytes::<C>::default();
        let mut sign = 0;

        loop {
            rng.try_fill_bytes(&mut bytes)?;
            rng.try_fill_bytes(core::array::from_mut(&mut sign))?;
            if let Some(point) = Self::decompress(&bytes, Choice::from(sign & 1)).into_option() {
                return Ok(point);
            }
        }
    }
}

impl<C> AffineCoordinates for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type FieldRepr = FieldBytes<C>;

    fn from_coordinates(x: &Self::FieldRepr, y: &Self::FieldRepr) -> CtOption<Self> {
        C::FieldElement::from_repr(*y).and_then(|y| {
            C::FieldElement::from_repr(*x).and_then(|x| {
                let lhs = y * &y;
                let rhs = x * &x * &x + &(C::EQUATION_A * &x) + &C::EQUATION_B;
                CtOption::new(Self { x, y, infinity: 0 }, lhs.ct_eq(&rhs))
            })
        })
    }

    fn x(&self) -> FieldBytes<C> {
        self.x.to_repr()
    }

    fn y(&self) -> FieldBytes<C> {
        self.y.to_repr()
    }

    fn x_is_odd(&self) -> Choice {
        self.x.is_odd()
    }

    fn y_is_odd(&self) -> Choice {
        self.y.is_odd()
    }
}

impl<C> ConditionallySelectable for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    #[inline(always)]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self {
            x: C::FieldElement::conditional_select(&a.x, &b.x, choice),
            y: C::FieldElement::conditional_select(&a.y, &b.y, choice),
            infinity: u8::conditional_select(&a.infinity, &b.infinity, choice),
        }
    }
}

impl<C> ConstantTimeEq for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn ct_eq(&self, other: &Self) -> Choice {
        self.x.ct_eq(&other.x) & self.y.ct_eq(&other.y) & self.infinity.ct_eq(&other.infinity)
    }
}

impl<C> ctutils::CtEq for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn ct_eq(&self, other: &Self) -> ctutils::Choice {
        ConstantTimeEq::ct_eq(self, other).into()
    }
}

impl<C> ctutils::CtSelect for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn ct_select(&self, other: &Self, choice: ctutils::Choice) -> Self {
        ConditionallySelectable::conditional_select(self, other, choice.into())
    }
}

impl<C> Default for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn default() -> Self {
        Self::IDENTITY
    }
}

impl<C> DefaultIsZeroes for AffinePoint<C> where C: PrimeCurveParams {}

impl<C> DecompressPoint<C> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn decompress(x_bytes: &FieldBytes<C>, y_is_odd: Choice) -> CtOption<Self> {
        C::FieldElement::from_repr(*x_bytes).and_then(|x| {
            let alpha = x * &x * &x + &(C::EQUATION_A * &x) + &C::EQUATION_B;
            let beta = alpha.sqrt();

            beta.map(|beta| {
                let y = C::FieldElement::conditional_select(
                    &-beta,
                    &beta,
                    beta.is_odd().ct_eq(&y_is_odd),
                );

                Self { x, y, infinity: 0 }
            })
        })
    }
}

impl<C> DecompactPoint<C> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn decompact(x_bytes: &FieldBytes<C>) -> CtOption<Self> {
        Self::decompress(x_bytes, Choice::from(0)).map(AffinePoint::to_compact)
    }
}

impl<C> Eq for AffinePoint<C> where C: PrimeCurveParams {}

impl<C> FromSec1Point<C> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    /// Attempts to parse the given [`Sec1Point`] as an SEC1-encoded
    /// [`AffinePoint`].
    ///
    /// # Returns
    ///
    /// `None` value if `encoded_point` is not on the secp384r1 curve.
    fn from_sec1_point(encoded_point: &Sec1Point<C>) -> ctutils::CtOption<Self> {
        match encoded_point.coordinates() {
            sec1::Coordinates::Identity => ctutils::CtOption::some(Self::IDENTITY),
            sec1::Coordinates::Compact { x } => Self::decompact(x).into(),
            sec1::Coordinates::Compressed { x, y_is_odd } => {
                Self::decompress(x, Choice::from(u8::from(y_is_odd))).into()
            }
            sec1::Coordinates::Uncompressed { x, y } => Self::from_coordinates(x, y).into(),
        }
    }
}

impl<C> From<NonIdentity<AffinePoint<C>>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn from(affine: NonIdentity<AffinePoint<C>>) -> Self {
        affine.to_point()
    }
}

impl<C> From<ProjectivePoint<C>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn from(p: ProjectivePoint<C>) -> AffinePoint<C> {
        p.to_affine()
    }
}

impl<C> From<&ProjectivePoint<C>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn from(p: &ProjectivePoint<C>) -> AffinePoint<C> {
        p.to_affine()
    }
}

impl<C> From<PublicKey<C>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn from(public_key: PublicKey<C>) -> AffinePoint<C> {
        *public_key.as_affine()
    }
}

impl<C> From<&PublicKey<C>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn from(public_key: &PublicKey<C>) -> AffinePoint<C> {
        AffinePoint::from(*public_key)
    }
}

impl<C> From<AffinePoint<C>> for Sec1Point<C>
where
    C: PrimeCurveParams,
{
    fn from(affine: AffinePoint<C>) -> Sec1Point<C> {
        affine.to_sec1_point(false)
    }
}

impl<C> Generate for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn try_generate_from_rng<R: TryCryptoRng + ?Sized>(
        rng: &mut R,
    ) -> core::result::Result<Self, R::Error> {
        Self::try_random(rng)
    }
}

impl<C> GroupEncoding for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type Repr = CompressedPoint<C>;

    /// NOTE: not constant-time with respect to identity point
    fn from_bytes(bytes: &Self::Repr) -> CtOption<Self> {
        Sec1Point::<C>::from_bytes(bytes)
            .map_or_else(
                |_| {
                    // SEC1 identity encoding is technically 1-byte 0x00, but the
                    // `GroupEncoding` API requires a fixed-width `Repr`
                    let is_identity =
                        ctutils::CtEq::ct_eq(bytes.as_slice(), Self::Repr::default().as_slice());
                    ctutils::CtOption::new(Sec1Point::<C>::identity(), is_identity)
                },
                ctutils::CtOption::some,
            )
            .and_then(|point| Self::from_sec1_point(&point))
            .into()
    }

    fn from_bytes_unchecked(bytes: &Self::Repr) -> CtOption<Self> {
        // No unchecked conversion possible for compressed points
        Self::from_bytes(bytes)
    }

    fn to_bytes(&self) -> Self::Repr {
        let encoded = self.to_sec1_point(true);
        let mut result = CompressedPoint::<C>::default();
        result[..encoded.len()].copy_from_slice(encoded.as_bytes());
        result
    }
}

impl<C> PartialEq for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn eq(&self, other: &Self) -> bool {
        self.ct_eq(other).into()
    }
}

impl<C> CurveAffine for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type Curve = ProjectivePoint<C>;
    type Scalar = Scalar<C>;

    fn identity() -> AffinePoint<C> {
        Self::IDENTITY
    }

    fn generator() -> AffinePoint<C> {
        Self::GENERATOR
    }

    fn is_identity(&self) -> Choice {
        self.is_identity()
    }

    fn to_curve(&self) -> ProjectivePoint<C> {
        ProjectivePoint::from(*self)
    }
}

impl<C> ToCompactSec1Point<C> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    /// Serialize this value as a  SEC1 compact [`Sec1Point`]
    fn to_compact_encoded_point(&self) -> ctutils::CtOption<Sec1Point<C>> {
        let point = self.to_compact();

        let mut bytes = CompressedPoint::<C>::default();
        bytes[0] = sec1::Tag::Compact.into();
        bytes[1..].copy_from_slice(&point.x.to_repr());

        let encoded = Sec1Point::<C>::from_bytes(bytes);
        let is_some =
            ctutils::CtEq::ct_eq(point.y.to_repr().as_slice(), self.y.to_repr().as_slice());
        ctutils::CtOption::new(encoded.unwrap_or_default(), is_some)
    }
}

impl<C> ToSec1Point<C> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn to_sec1_point(&self, compress: bool) -> Sec1Point<C> {
        Sec1Point::<C>::ct_select(
            &Sec1Point::<C>::from_affine_coordinates(
                &self.x.to_repr(),
                &self.y.to_repr(),
                compress,
            ),
            &Sec1Point::<C>::identity(),
            self.is_identity().into(),
        )
    }
}

/// The constant-time alternative is available at [`NonIdentity::new()`].
impl<C> TryFrom<AffinePoint<C>> for NonIdentity<AffinePoint<C>>
where
    C: PrimeCurveParams,
{
    type Error = Error;

    fn try_from(affine_point: AffinePoint<C>) -> Result<Self> {
        NonIdentity::new(affine_point).into_option().ok_or(Error)
    }
}

impl<C> TryFrom<Sec1Point<C>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type Error = Error;

    fn try_from(point: Sec1Point<C>) -> Result<AffinePoint<C>> {
        AffinePoint::try_from(&point)
    }
}

impl<C> TryFrom<&Sec1Point<C>> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type Error = Error;

    fn try_from(point: &Sec1Point<C>) -> Result<AffinePoint<C>> {
        Option::from(AffinePoint::<C>::from_sec1_point(point)).ok_or(Error)
    }
}

impl<C> TryFrom<AffinePoint<C>> for PublicKey<C>
where
    C: PrimeCurveParams,
{
    type Error = Error;

    fn try_from(affine_point: AffinePoint<C>) -> Result<PublicKey<C>> {
        PublicKey::from_affine(affine_point)
    }
}

impl<C> TryFrom<&AffinePoint<C>> for PublicKey<C>
where
    C: PrimeCurveParams,
{
    type Error = Error;

    fn try_from(affine_point: &AffinePoint<C>) -> Result<PublicKey<C>> {
        PublicKey::<C>::try_from(*affine_point)
    }
}

//
// Arithmetic trait impls
//

impl<C, S> Mul<S> for AffinePoint<C>
where
    C: PrimeCurveParams,
    S: Borrow<Scalar<C>>,
    ProjectivePoint<C>: Double,
{
    type Output = ProjectivePoint<C>;

    #[inline]
    fn mul(self, scalar: S) -> ProjectivePoint<C> {
        ProjectivePoint::<C>::from(self) * scalar
    }
}

impl<C, S> Mul<S> for &AffinePoint<C>
where
    C: PrimeCurveParams,
    S: Borrow<Scalar<C>>,
    ProjectivePoint<C>: Double,
{
    type Output = ProjectivePoint<C>;

    #[inline]
    fn mul(self, scalar: S) -> ProjectivePoint<C> {
        ProjectivePoint::<C>::from(self) * scalar
    }
}

impl<C, S> MulVartime<S> for AffinePoint<C>
where
    C: PrimeCurveParams,
    S: Borrow<Scalar<C>>,
    ProjectivePoint<C>: Double,
{
    #[inline]
    fn mul_vartime(self, scalar: S) -> ProjectivePoint<C> {
        ProjectivePoint::<C>::from(self).mul_vartime(scalar)
    }
}

impl<C, S> MulVartime<S> for &AffinePoint<C>
where
    C: PrimeCurveParams,
    S: Borrow<Scalar<C>>,
    ProjectivePoint<C>: Double,
{
    #[inline]
    fn mul_vartime(self, scalar: S) -> ProjectivePoint<C> {
        ProjectivePoint::<C>::from(self).mul_vartime(scalar)
    }
}

impl<C> Neg for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type Output = Self;

    #[inline]
    fn neg(self) -> Self {
        AffinePoint {
            x: self.x,
            y: -self.y,
            infinity: self.infinity,
        }
    }
}

impl<C> Neg for &AffinePoint<C>
where
    C: PrimeCurveParams,
{
    type Output = AffinePoint<C>;

    #[inline]
    fn neg(self) -> AffinePoint<C> {
        -(*self)
    }
}

//
// serde support
//

#[cfg(feature = "serde")]
impl<C> Serialize for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn serialize<S>(&self, serializer: S) -> core::result::Result<S::Ok, S::Error>
    where
        S: ser::Serializer,
    {
        self.to_sec1_point(true).serialize(serializer)
    }
}

#[cfg(feature = "serde")]
impl<'de, C> Deserialize<'de> for AffinePoint<C>
where
    C: PrimeCurveParams,
{
    fn deserialize<D>(deserializer: D) -> core::result::Result<Self, D::Error>
    where
        D: de::Deserializer<'de>,
    {
        Sec1Point::<C>::deserialize(deserializer)?
            .try_into()
            .map_err(de::Error::custom)
    }
}
