//! Elliptic curve domain parameters.

#![allow(unreachable_patterns)]
#![allow(deprecated)]

use alloc::vec::Vec;
use core::fmt::Debug;
use rand_core::{CryptoRng, RngCore};
use subtle::{Choice, CtOption};
use zeroize::Zeroize;

use crate::error::AgreementError;

// Conditional imports for elliptic-curve common traits
#[cfg(any(feature = "secp256k1", feature = "secp256r1"))]
use elliptic_curve::{
    bigint::{Encoding, U256},
    group::GroupEncoding,
    sec1::{FromEncodedPoint, ToEncodedPoint},
    PrimeField, // Required for `from_repr`
};

#[cfg(feature = "curve25519")]
use curve25519_dalek::{
    edwards::{CompressedEdwardsY, EdwardsPoint},
    scalar::Scalar as DalekScalar,
    traits::{IsIdentity, Identity},
};

#[cfg(feature = "secp256k1")]
use k256::{
    AffinePoint as K256Affine,
    ProjectivePoint as K256Projective,
    Scalar as K256Scalar,
    FieldBytes as K256FieldBytes,
};

#[cfg(feature = "secp256r1")]
use p256::{
    AffinePoint as P256Affine,
    ProjectivePoint as P256Projective,
    Scalar as P256Scalar,
    FieldBytes as P256FieldBytes,
};

// -----------------------------------------------------------------------------
// Curve trait
// -----------------------------------------------------------------------------

/// Core trait defining the elliptic curve arithmetic required by ECMQV.
///
/// This trait abstracts the underlying prime-order curve operations. Implementations
/// are provided for `Curve25519`, `secp256k1`, and `secp256r1`.
pub trait Curve: Debug + Send + Sync + Clone + Copy + 'static {
    /// Type representing a field element (integer modulo the field prime).
    type FieldElement: Clone + Copy + Debug;
    /// Type representing a scalar (secret exponent) in the curve's scalar field.
    type Scalar: Clone + Copy + Debug + PartialEq + Eq + Zeroize;
    /// Type representing an affine point (public key) on the curve.
    type PublicKey: Clone + Copy + Debug + PartialEq + Eq;

    /// Returns the curve identifier.
    fn curve_id() -> CurveId;
    /// Returns the field size in bits.
    fn field_bits() -> usize;
    /// Returns half the field size, rounded up (used for association value computation).
    fn half_len() -> usize {
        (Self::field_bits() + 1) / 2
    }
    /// Returns the cofactor (h) of the curve.
    fn cofactor() -> u32;
    /// Returns the order (n) of the curve's prime-order subgroup.
    fn order() -> Self::Scalar;
    /// Generates a random scalar in the range [1, n-1].
    fn random_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Scalar;
    /// Converts a big-endian byte array to a scalar, returning `None` if the value is not in the field.
    fn scalar_from_be_bytes(bytes: &[u8; 32]) -> CtOption<Self::Scalar>;
    /// Serializes a scalar to big-endian bytes.
    fn scalar_to_bytes(scalar: &Self::Scalar) -> Vec<u8>;
    /// Computes the association value `\bar{x}` for a public key (used in ECMQV).
    fn associate_value(point: &Self::PublicKey) -> Self::Scalar;
    /// Scalar multiplication by the curve's base point.
    fn mul_base(scalar: &Self::Scalar) -> Self::PublicKey;
    /// Scalar multiplication of a point by a scalar.
    fn mul_scalar(point: &Self::PublicKey, scalar: &Self::Scalar) -> Self::PublicKey;
    /// Point addition.
    fn add_points(a: &Self::PublicKey, b: &Self::PublicKey) -> Self::PublicKey;
    /// Returns `true` if the point is the point at infinity.
    fn is_identity(point: &Self::PublicKey) -> bool;
    /// Validates a public key (not infinity, on curve, in the correct subgroup).
    fn validate_public_key(point: &Self::PublicKey) -> Result<(), AgreementError>;
    /// Serializes a public key to compressed bytes.
    fn public_key_to_bytes(point: &Self::PublicKey) -> Vec<u8>;
    /// Deserializes a public key from compressed bytes.
    fn public_key_from_bytes(bytes: &[u8]) -> CtOption<Self::PublicKey>;

    /// Returns the zero scalar.
    fn zero_scalar() -> Self::Scalar;
    /// Returns the identity point (point at infinity).
    fn identity_point() -> Self::PublicKey;
    /// Returns the cofactor as a scalar.
    fn cofactor_scalar() -> Self::Scalar;
}

// =============================================================================
// Curve25519 implementation
// =============================================================================

/// Concrete implementation for Curve25519 (Edwards25519).
#[cfg(feature = "curve25519")]
#[derive(Debug, Clone, Copy)]
pub struct Curve25519;

#[cfg(feature = "curve25519")]
impl Curve for Curve25519 {
    type FieldElement = u64;
    type Scalar = DalekScalar;
    type PublicKey = EdwardsPoint;

    fn curve_id() -> CurveId {
        CurveId::Curve25519
    }
    fn field_bits() -> usize {
        255
    }
    fn cofactor() -> u32 {
        8
    }

    fn order() -> Self::Scalar {
        DalekScalar::from_bytes_mod_order([
            0xed, 0xd3, 0xf5, 0x5c, 0x1a, 0x63, 0x12, 0x58,
            0xd6, 0x9c, 0xf7, 0xa2, 0xde, 0xf9, 0xde, 0x14,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10,
        ])
    }

    fn random_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Scalar {
        let mut bytes = [0u8; 32];
        rng.fill_bytes(&mut bytes);
        DalekScalar::from_bytes_mod_order(bytes)
    }

    fn scalar_from_be_bytes(bytes: &[u8; 32]) -> CtOption<Self::Scalar> {
        let mut le = [0u8; 32];
        for i in 0..32 {
            le[i] = bytes[31 - i];
        }
        CtOption::new(DalekScalar::from_bytes_mod_order(le), Choice::from(1))
    }

    fn scalar_to_bytes(scalar: &Self::Scalar) -> Vec<u8> {
        scalar.as_bytes().to_vec()
    }

    fn associate_value(point: &Self::PublicKey) -> Self::Scalar {
        let comp = point.compress();
        let x = comp.as_bytes();
        let mut low = [0u8; 16];
        low.copy_from_slice(&x[0..16]);
        let low_val = u128::from_le_bytes(low);
        let mut two = [0u8; 32];
        two[16] = 1;
        DalekScalar::from(low_val) + DalekScalar::from_bytes_mod_order(two)
    }

    fn mul_base(scalar: &Self::Scalar) -> Self::PublicKey {
        EdwardsPoint::mul_base(scalar)
    }
    fn mul_scalar(point: &Self::PublicKey, scalar: &Self::Scalar) -> Self::PublicKey {
        scalar * point
    }
    fn add_points(a: &Self::PublicKey, b: &Self::PublicKey) -> Self::PublicKey {
        a + b
    }
    fn is_identity(point: &Self::PublicKey) -> bool {
        point.is_identity()
    }

    fn validate_public_key(point: &Self::PublicKey) -> Result<(), AgreementError> {
        if point.is_identity() {
            return Err(AgreementError::InvalidPublicKey);
        }
        if point.compress().decompress().is_none() {
            return Err(AgreementError::InvalidPublicKey);
        }
        let c = DalekScalar::from(Self::cofactor()) * point;
        if c.is_identity() {
            return Err(AgreementError::InvalidPublicKey);
        }
        Ok(())
    }

    fn public_key_to_bytes(point: &Self::PublicKey) -> Vec<u8> {
        point.compress().as_bytes().to_vec()
    }

    fn public_key_from_bytes(bytes: &[u8]) -> CtOption<Self::PublicKey> {
        if bytes.len() != 32 {
            return CtOption::new(EdwardsPoint::identity(), Choice::from(0));
        }
        let mut arr = [0u8; 32];
        arr.copy_from_slice(bytes);
        match CompressedEdwardsY(arr).decompress() {
            Some(p) => CtOption::new(p, Choice::from(1)),
            None => CtOption::new(EdwardsPoint::identity(), Choice::from(0)),
        }
    }

    fn zero_scalar() -> Self::Scalar {
        DalekScalar::from(0u64)
    }
    fn identity_point() -> Self::PublicKey {
        EdwardsPoint::identity()
    }
    fn cofactor_scalar() -> Self::Scalar {
        DalekScalar::from(Self::cofactor())
    }
}

// =============================================================================
// Secp256k1 implementation
// =============================================================================

/// Concrete implementation for the secp256k1 (Koblitz) curve.
#[cfg(feature = "secp256k1")]
#[derive(Debug, Clone, Copy)]
pub struct Secp256k1;

#[cfg(feature = "secp256k1")]
impl Curve for Secp256k1 {
    type FieldElement = U256;
    type Scalar = K256Scalar;
    type PublicKey = K256Affine;

    fn curve_id() -> CurveId {
        CurveId::Secp256k1
    }
    fn field_bits() -> usize {
        256
    }
    fn cofactor() -> u32 {
        1
    }

    fn order() -> Self::Scalar {
        const ORDER: [u8; 32] = [
            0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
            0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFE,
            0xBA, 0xAE, 0xDC, 0xE6, 0xAF, 0x48, 0xA0, 0x3B,
            0xBF, 0xD2, 0x5E, 0x8C, 0xD0, 0x36, 0x41, 0x41,
        ];
        // SAFETY: val_bytes is always within the field, so unwrap is safe.
        K256Scalar::from_repr(K256FieldBytes::from(ORDER)).unwrap()
    }

    fn random_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Scalar {
        loop {
            let mut bytes = [0u8; 32];
            rng.fill_bytes(&mut bytes);
            if let Some(s) = K256Scalar::from_repr(K256FieldBytes::from(bytes)).into_option() {
                return s;
            }
        }
    }

    fn scalar_from_be_bytes(bytes: &[u8; 32]) -> CtOption<Self::Scalar> {
        K256Scalar::from_repr(K256FieldBytes::from(*bytes))
    }

    fn scalar_to_bytes(scalar: &Self::Scalar) -> Vec<u8> {
        scalar.to_bytes().as_slice().to_vec()
    }

    fn associate_value(point: &Self::PublicKey) -> Self::Scalar {
        let x_bytes = point.to_bytes();
        let x_raw = if x_bytes[0] == 0x04 {
            let mut x = [0u8; 32];
            x.copy_from_slice(&x_bytes[1..33]);
            U256::from_be_slice(&x)
        } else {
            let mut x = [0u8; 32];
            x.copy_from_slice(&x_bytes[1..33]);
            U256::from_be_slice(&x)
        };
        let half = Self::half_len();
        let mod2 = U256::from(1u64) << half;
        let x_mod = x_raw.wrapping_rem(&mod2.into());
        let val = x_mod.wrapping_add(&mod2);
        let val_bytes = val.to_be_bytes();
        K256Scalar::from_repr(K256FieldBytes::from(val_bytes)).unwrap()
    }

    fn mul_base(scalar: &Self::Scalar) -> Self::PublicKey {
        (K256Projective::GENERATOR * scalar).to_affine()
    }
    fn mul_scalar(point: &Self::PublicKey, scalar: &Self::Scalar) -> Self::PublicKey {
        (K256Projective::from(*point) * scalar).to_affine()
    }
    fn add_points(a: &Self::PublicKey, b: &Self::PublicKey) -> Self::PublicKey {
        (K256Projective::from(*a) + K256Projective::from(*b)).to_affine()
    }
    fn is_identity(point: &Self::PublicKey) -> bool {
        use k256::elliptic_curve::group::prime::PrimeCurveAffine;
        bool::from(point.is_identity())
    }

    fn validate_public_key(point: &Self::PublicKey) -> Result<(), AgreementError> {
        use k256::elliptic_curve::group::prime::PrimeCurveAffine;
        if bool::from(point.is_identity()) {
            return Err(AgreementError::InvalidPublicKey);
        }
        let encoded = point.to_encoded_point(true);
        let p = K256Affine::from_encoded_point(&encoded);
        if bool::from(p.is_some()) {
            return Ok(());
        }
        Err(AgreementError::InvalidPublicKey)
    }

    fn public_key_to_bytes(point: &Self::PublicKey) -> Vec<u8> {
        point.to_encoded_point(true).as_bytes().to_vec()
    }

    fn public_key_from_bytes(bytes: &[u8]) -> CtOption<Self::PublicKey> {
        if let Ok(enc) = k256::EncodedPoint::from_bytes(bytes) {
            return K256Affine::from_encoded_point(&enc);
        }
        CtOption::new(K256Affine::IDENTITY, Choice::from(0))
    }

    fn zero_scalar() -> Self::Scalar {
        K256Scalar::from(0u64)
    }
    fn identity_point() -> Self::PublicKey {
        K256Affine::IDENTITY
    }
    fn cofactor_scalar() -> Self::Scalar {
        K256Scalar::from(Self::cofactor())
    }
}

// =============================================================================
// Secp256r1 (NIST P-256) implementation
// =============================================================================

/// Concrete implementation for the NIST P-256 (secp256r1) curve.
#[cfg(feature = "secp256r1")]
#[derive(Debug, Clone, Copy)]
pub struct Secp256r1;

#[cfg(feature = "secp256r1")]
impl Curve for Secp256r1 {
    type FieldElement = U256;
    type Scalar = P256Scalar;
    type PublicKey = P256Affine;

    fn curve_id() -> CurveId {
        CurveId::Secp256r1
    }
    fn field_bits() -> usize {
        256
    }
    fn cofactor() -> u32 {
        1
    }

    fn order() -> Self::Scalar {
        const ORDER: [u8; 32] = [
            0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00,
            0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
            0xBC, 0xE6, 0xFA, 0xAD, 0xA7, 0x17, 0x9E, 0x84,
            0xF3, 0xB9, 0xCA, 0xC2, 0xFC, 0x63, 0x25, 0x51,
        ];
        P256Scalar::from_repr(P256FieldBytes::from(ORDER)).unwrap()
    }

    fn random_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Scalar {
        loop {
            let mut bytes = [0u8; 32];
            rng.fill_bytes(&mut bytes);
            if let Some(s) = P256Scalar::from_repr(P256FieldBytes::from(bytes)).into_option() {
                return s;
            }
        }
    }

    fn scalar_from_be_bytes(bytes: &[u8; 32]) -> CtOption<Self::Scalar> {
        P256Scalar::from_repr(P256FieldBytes::from(*bytes))
    }

    fn scalar_to_bytes(scalar: &Self::Scalar) -> Vec<u8> {
        scalar.to_bytes().as_slice().to_vec()
    }

    fn associate_value(point: &Self::PublicKey) -> Self::Scalar {
        let x_bytes = point.to_bytes();
        let x_raw = if x_bytes[0] == 0x04 {
            let mut x = [0u8; 32];
            x.copy_from_slice(&x_bytes[1..33]);
            U256::from_be_slice(&x)
        } else {
            let mut x = [0u8; 32];
            x.copy_from_slice(&x_bytes[1..33]);
            U256::from_be_slice(&x)
        };
        let half = Self::half_len();
        let mod2 = U256::from(1u64) << half;
        let x_mod = x_raw.wrapping_rem(&mod2.into());
        let val = x_mod.wrapping_add(&mod2);
        let val_bytes = val.to_be_bytes();
        P256Scalar::from_repr(P256FieldBytes::from(val_bytes)).unwrap()
    }

    fn mul_base(scalar: &Self::Scalar) -> Self::PublicKey {
        (P256Projective::GENERATOR * scalar).to_affine()
    }
    fn mul_scalar(point: &Self::PublicKey, scalar: &Self::Scalar) -> Self::PublicKey {
        (P256Projective::from(*point) * scalar).to_affine()
    }
    fn add_points(a: &Self::PublicKey, b: &Self::PublicKey) -> Self::PublicKey {
        (P256Projective::from(*a) + P256Projective::from(*b)).to_affine()
    }
    fn is_identity(point: &Self::PublicKey) -> bool {
        bool::from(point.is_identity())
    }

    fn validate_public_key(point: &Self::PublicKey) -> Result<(), AgreementError> {
        if bool::from(point.is_identity()) {
            return Err(AgreementError::InvalidPublicKey);
        }
        let encoded = point.to_encoded_point(true);
        let p = P256Affine::from_encoded_point(&encoded);
        if bool::from(p.is_some()) {
            return Ok(());
        }
        Err(AgreementError::InvalidPublicKey)
    }

    fn public_key_to_bytes(point: &Self::PublicKey) -> Vec<u8> {
        point.to_encoded_point(true).as_bytes().to_vec()
    }

    fn public_key_from_bytes(bytes: &[u8]) -> CtOption<Self::PublicKey> {
        if let Ok(enc) = p256::EncodedPoint::from_bytes(bytes) {
            return P256Affine::from_encoded_point(&enc);
        }
        CtOption::new(P256Affine::IDENTITY, Choice::from(0))
    }

    fn zero_scalar() -> Self::Scalar {
        P256Scalar::from(0u64)
    }
    fn identity_point() -> Self::PublicKey {
        P256Affine::IDENTITY
    }
    fn cofactor_scalar() -> Self::Scalar {
        P256Scalar::from(Self::cofactor())
    }
}

// =============================================================================
// CurveId & CurveParameters
// =============================================================================

/// Supported elliptic curve identifiers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum CurveId {
    /// Edwards25519 curve (Curve25519).
    Curve25519,
    /// Koblitz curve secp256k1.
    Secp256k1,
    /// NIST P‑256 curve secp256r1.
    Secp256r1,
}

/// Curve parameter set, used to select a curve at runtime.
#[derive(Debug, Clone, Copy)]
pub enum CurveParameters {
    /// Curve25519
    Curve25519,
    /// secp256k1
    #[cfg(feature = "secp256k1")]
    Secp256k1,
    /// secp256r1
    #[cfg(feature = "secp256r1")]
    Secp256r1,
}

impl CurveParameters {
    /// Returns the curve identifier for this parameter set.
    pub fn curve_id(&self) -> CurveId {
        match self {
            Self::Curve25519 => CurveId::Curve25519,
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1 => CurveId::Secp256k1,
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1 => CurveId::Secp256r1,
        }
    }

    /// Returns half the field size in bits for this curve.
    pub fn half_len(&self) -> usize {
        match self {
            Self::Curve25519 => Curve25519::half_len(),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1 => Secp256k1::half_len(),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1 => Secp256r1::half_len(),
        }
    }

    /// Computes the association value (scalar) from the byte representation of a public key.
    ///
    /// # Errors
    /// Returns `InvalidPublicKey` if the bytes do not decode to a valid point.
    pub fn associate_value_from_bytes(&self, bytes: &[u8]) -> Result<Vec<u8>, AgreementError> {
        match self {
            Self::Curve25519 => {
                #[cfg(feature = "curve25519")]
                {
                    let p = Curve25519::public_key_from_bytes(bytes);
                    if p.is_none().into() {
                        return Err(AgreementError::InvalidPublicKey);
                    }
                    let s = Curve25519::associate_value(&p.unwrap());
                    Ok(Curve25519::scalar_to_bytes(&s))
                }
                #[cfg(not(feature = "curve25519"))]
                Err(AgreementError::UnsupportedCurve)
            }
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1 => {
                let p = Secp256k1::public_key_from_bytes(bytes);
                if p.is_none().into() {
                    return Err(AgreementError::InvalidPublicKey);
                }
                let s = Secp256k1::associate_value(&p.unwrap());
                Ok(Secp256k1::scalar_to_bytes(&s))
            }
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1 => {
                let p = Secp256r1::public_key_from_bytes(bytes);
                if p.is_none().into() {
                    return Err(AgreementError::InvalidPublicKey);
                }
                let s = Secp256r1::associate_value(&p.unwrap());
                Ok(Secp256r1::scalar_to_bytes(&s))
            }
        }
    }
}