//! Key types: ephemeral keys, static keys, and public keys.
//!
//! This module provides the key structures used in ECMQV key exchange.
//! All sensitive fields (secret scalars) implement `ZeroizeOnDrop` automatically
//! via the `Scalar` type from the underlying curve libraries.
#![allow(unreachable_patterns)]
#![allow(deprecated)]
use alloc::vec::Vec;
use rand_core::{CryptoRng, RngCore};
use zeroize::Zeroize;


#[cfg(feature = "curve25519")]
use curve25519_dalek::edwards::EdwardsPoint;

#[cfg(feature = "secp256k1")]
use k256::AffinePoint as K256Affine;

#[cfg(feature = "secp256r1")]
use p256::AffinePoint as P256Affine;

use crate::error::AgreementError;
use crate::params::{Curve, CurveParameters, CurveId};

// -----------------------------------------------------------------------------
// Internal generic key container (for a specific curve)
// -----------------------------------------------------------------------------

/// Internal ephemeral key for a concrete curve `C`.
///
/// This holds the secret scalar and the corresponding public point.
/// The secret is automatically zeroized on drop.
pub struct EphemeralKeyInner<C: Curve> {
    secret: C::Scalar,
    public: C::PublicKey,
}

impl<C: Curve> EphemeralKeyInner<C> {
    /// Generates a new ephemeral key pair using a cryptographically secure RNG.
    pub fn generate<R: RngCore + CryptoRng>(rng: &mut R) -> Self
    where
        C::Scalar: Zeroize,
    {
        let secret = C::random_scalar(rng);
        let public = C::mul_base(&secret);
        Self { secret, public }
    }

    /// Accesses the secret scalar (internal use only).
    pub(crate) fn secret(&self) -> &C::Scalar {
        &self.secret
    }

    /// Accesses the public key (internal use only).
    pub(crate) fn public(&self) -> &C::PublicKey {
        &self.public
    }
}

/// Internal static key for a concrete curve `C`.
///
/// This holds a long‑term secret scalar and its public point.
/// The secret is automatically zeroized on drop.
pub struct StaticKeyInner<C: Curve> {
    secret: C::Scalar,
    public: C::PublicKey,
}

impl<C: Curve> StaticKeyInner<C> {
    /// Derives a static key from a 32‑byte seed.
    ///
    /// If the seed is not a valid scalar, a random scalar is generated instead.
    pub fn from_seed<R: RngCore + CryptoRng>(seed: &[u8; 32], rng: &mut R) -> Self
    where
        C::Scalar: Zeroize + From<u64>,
    {
        let secret = C::scalar_from_be_bytes(seed);
        let secret = if secret.is_none().into() {
            C::random_scalar(rng)
        } else {
            secret.unwrap()
        };
        let public = C::mul_base(&secret);
        Self { secret, public }
    }

    /// Deserializes a static key from its secret bytes (big‑endian).
    ///
    /// # Errors
    /// Returns `InvalidPublicKey` if the bytes are not a valid scalar.
    pub fn from_bytes(secret_bytes: &[u8]) -> Result<Self, AgreementError>
    where
        C::Scalar: Zeroize + From<u64>,
    {
        if secret_bytes.len() != 32 {
            return Err(AgreementError::InvalidPublicKey);
        }
        let mut arr = [0u8; 32];
        arr.copy_from_slice(secret_bytes);
        let secret = C::scalar_from_be_bytes(&arr);
        if secret.is_none().into() {
            return Err(AgreementError::InvalidPublicKey);
        }
        let secret = secret.unwrap();
        let public = C::mul_base(&secret);
        Ok(Self { secret, public })
    }

    /// Accesses the secret scalar (internal use only).
    pub(crate) fn secret(&self) -> &C::Scalar {
        &self.secret
    }

    /// Accesses the public key (internal use only).
    pub(crate) fn public(&self) -> &C::PublicKey {
        &self.public
    }
}

// -----------------------------------------------------------------------------
// PublicKey enum (wrapper over curve‑specific public keys)
// -----------------------------------------------------------------------------

/// A public key (point on the elliptic curve) used in ECMQV.
///
/// This enum holds a point on one of the supported curves.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum PublicKey {
    /// Curve25519 public key (compressed Edwards point).
    #[cfg(feature = "curve25519")]
    Curve25519(EdwardsPoint),
    /// secp256k1 public key (affine point).
    #[cfg(feature = "secp256k1")]
    Secp256k1(K256Affine),
    /// secp256r1 (NIST P‑256) public key.
    #[cfg(feature = "secp256r1")]
    Secp256r1(P256Affine),
}

impl PublicKey {
    /// Validates the public key according to SEC 1: not identity, on curve, in the correct subgroup.
    ///
    /// # Errors
    /// Returns `InvalidPublicKey` if the key fails validation.
    pub fn validate(&self) -> Result<(), AgreementError> {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(point) => crate::params::Curve25519::validate_public_key(point),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(point) => crate::params::Secp256k1::validate_public_key(point),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(point) => crate::params::Secp256r1::validate_public_key(point),
            #[allow(unreachable_patterns)]
            _ => Err(AgreementError::UnsupportedCurve),
        }
    }

    /// Serializes the public key to bytes in curve‑specific compressed format.
    pub fn to_bytes(&self) -> alloc::vec::Vec<u8> {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(point) => crate::params::Curve25519::public_key_to_bytes(point),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(point) => crate::params::Secp256k1::public_key_to_bytes(point),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(point) => crate::params::Secp256r1::public_key_to_bytes(point),
            #[allow(unreachable_patterns)]
            _ => alloc::vec![],
        }
    }

    /// Deserializes a public key from bytes given the expected curve.
    ///
    /// # Errors
    /// Returns `InvalidPublicKey` if the bytes do not decode to a valid point.
    /// Returns `UnsupportedCurve` if the curve is not enabled.
    pub fn from_bytes(bytes: &[u8], curve: CurveParameters) -> Result<Self, AgreementError> {
        match curve {
            CurveParameters::Curve25519 => {
                #[cfg(feature = "curve25519")]
                {
                    let point = crate::params::Curve25519::public_key_from_bytes(bytes);
                    if point.is_none().into() {
                        return Err(AgreementError::InvalidPublicKey);
                    }
                    Ok(PublicKey::Curve25519(point.unwrap()))
                }
                #[cfg(not(feature = "curve25519"))]
                Err(AgreementError::UnsupportedCurve)
            }
            #[cfg(feature = "secp256k1")]
            CurveParameters::Secp256k1 => {
                let point = crate::params::Secp256k1::public_key_from_bytes(bytes);
                if point.is_none().into() {
                    return Err(AgreementError::InvalidPublicKey);
                }
                Ok(PublicKey::Secp256k1(point.unwrap()))
            }
            #[cfg(feature = "secp256r1")]
            CurveParameters::Secp256r1 => {
                let point = crate::params::Secp256r1::public_key_from_bytes(bytes);
                if point.is_none().into() {
                    return Err(AgreementError::InvalidPublicKey);
                }
                Ok(PublicKey::Secp256r1(point.unwrap()))
            }
            _ => Err(AgreementError::UnsupportedCurve),
        }
    }

    /// Returns the curve identifier of this key.
    pub fn curve_id(&self) -> CurveId {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(_) => CurveId::Curve25519,
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(_) => CurveId::Secp256k1,
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(_) => CurveId::Secp256r1,
            #[allow(unreachable_patterns)]
            _ => unreachable!(),
        }
    }
}

// -----------------------------------------------------------------------------
// EphemeralKey enum
// -----------------------------------------------------------------------------

/// An ephemeral (short‑lived) key pair for a single session.
///
/// This key MUST be discarded after the key exchange and MUST NOT be reused.
#[non_exhaustive]
pub enum EphemeralKey {
    /// Curve25519 ephemeral key pair.
    #[cfg(feature = "curve25519")]
    Curve25519(EphemeralKeyInner<crate::params::Curve25519>),
    /// secp256k1 ephemeral key pair.
    #[cfg(feature = "secp256k1")]
    Secp256k1(EphemeralKeyInner<crate::params::Secp256k1>),
    /// secp256r1 ephemeral key pair.
    #[cfg(feature = "secp256r1")]
    Secp256r1(EphemeralKeyInner<crate::params::Secp256r1>),
}

impl EphemeralKey {
    /// Generates a new ephemeral key pair on the given curve.
    ///
    /// # Errors
    /// Returns `UnsupportedCurve` if the curve is not enabled in this build.
    pub fn generate<R: RngCore + CryptoRng>(
        rng: &mut R,
        curve: CurveParameters,
    ) -> Result<Self, AgreementError> {
        match curve {
            CurveParameters::Curve25519 => {
                #[cfg(feature = "curve25519")]
                {
                    Ok(EphemeralKey::Curve25519(EphemeralKeyInner::generate(rng)))
                }
                #[cfg(not(feature = "curve25519"))]
                Err(AgreementError::UnsupportedCurve)
            }
            #[cfg(feature = "secp256k1")]
            CurveParameters::Secp256k1 => Ok(EphemeralKey::Secp256k1(
                EphemeralKeyInner::generate(rng),
            )),
            #[cfg(feature = "secp256r1")]
            CurveParameters::Secp256r1 => {
                Ok(EphemeralKey::Secp256r1(EphemeralKeyInner::generate(rng)))
            }
            _ => Err(AgreementError::UnsupportedCurve),
        }
    }

    /// Returns the public key component.
    pub fn public(&self) -> PublicKey {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(inner) => PublicKey::Curve25519(*inner.public()),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(inner) => PublicKey::Secp256k1(*inner.public()),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(inner) => PublicKey::Secp256r1(*inner.public()),
            #[allow(unreachable_patterns)]
            _ => unreachable!(),
        }
    }

    /// Exports the secret scalar as bytes (for internal testing only).
    #[doc(hidden)]
    pub fn secret_scalar(&self) -> Vec<u8> {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(inner) => inner.secret().as_bytes().to_vec(),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(inner) => inner.secret().to_bytes().as_slice().to_vec(),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(inner) => inner.secret().to_bytes().as_slice().to_vec(),
            #[allow(unreachable_patterns)]
            _ => Vec::new(),
        }
    }

    /// Returns the curve identifier of this key.
    pub fn curve_id(&self) -> CurveId {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(_) => CurveId::Curve25519,
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(_) => CurveId::Secp256k1,
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(_) => CurveId::Secp256r1,
            #[allow(unreachable_patterns)]
            _ => unreachable!(),
        }
    }
}

// -----------------------------------------------------------------------------
// StaticKey enum
// -----------------------------------------------------------------------------

/// A static (long‑term) key pair, typically stored in a certificate.
///
/// This key can be reused across many sessions and should be stored securely.
#[non_exhaustive]
pub enum StaticKey {
    /// Curve25519 static key pair.
    #[cfg(feature = "curve25519")]
    Curve25519(StaticKeyInner<crate::params::Curve25519>),
    /// secp256k1 static key pair.
    #[cfg(feature = "secp256k1")]
    Secp256k1(StaticKeyInner<crate::params::Secp256k1>),
    /// secp256r1 static key pair.
    #[cfg(feature = "secp256r1")]
    Secp256r1(StaticKeyInner<crate::params::Secp256r1>),
}

impl StaticKey {
    /// Derives a static key from a 32‑byte seed.
    ///
    /// If the seed is not a valid scalar, a random scalar is generated.
    ///
    /// # Errors
    /// Returns `UnsupportedCurve` if the curve is not enabled.
    pub fn from_seed<R: RngCore + CryptoRng>(
        seed: &[u8; 32],
        rng: &mut R,
        curve: CurveParameters,
    ) -> Result<Self, AgreementError> {
        match curve {
            CurveParameters::Curve25519 => {
                #[cfg(feature = "curve25519")]
                {
                    Ok(StaticKey::Curve25519(StaticKeyInner::from_seed(seed, rng)))
                }
                #[cfg(not(feature = "curve25519"))]
                Err(AgreementError::UnsupportedCurve)
            }
            #[cfg(feature = "secp256k1")]
            CurveParameters::Secp256k1 => Ok(StaticKey::Secp256k1(StaticKeyInner::from_seed(
                seed, rng,
            ))),
            #[cfg(feature = "secp256r1")]
            CurveParameters::Secp256r1 => Ok(StaticKey::Secp256r1(StaticKeyInner::from_seed(
                seed, rng,
            ))),
            _ => Err(AgreementError::UnsupportedCurve),
        }
    }

    /// Deserializes a static key from its secret bytes.
    ///
    /// # Errors
    /// Returns `InvalidPublicKey` if the bytes are not a valid scalar.
    /// Returns `UnsupportedCurve` if the curve is not enabled.
    pub fn from_bytes(
        secret_bytes: &[u8],
        curve: CurveParameters,
    ) -> Result<Self, AgreementError> {
        match curve {
            CurveParameters::Curve25519 => {
                #[cfg(feature = "curve25519")]
                {
                    Ok(StaticKey::Curve25519(StaticKeyInner::from_bytes(secret_bytes)?))
                }
                #[cfg(not(feature = "curve25519"))]
                Err(AgreementError::UnsupportedCurve)
            }
            #[cfg(feature = "secp256k1")]
            CurveParameters::Secp256k1 => Ok(StaticKey::Secp256k1(StaticKeyInner::from_bytes(
                secret_bytes,
            )?)),
            #[cfg(feature = "secp256r1")]
            CurveParameters::Secp256r1 => Ok(StaticKey::Secp256r1(StaticKeyInner::from_bytes(
                secret_bytes,
            )?)),
            _ => Err(AgreementError::UnsupportedCurve),
        }
    }

    /// Returns the public key component.
    pub fn public(&self) -> PublicKey {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(inner) => PublicKey::Curve25519(*inner.public()),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(inner) => PublicKey::Secp256k1(*inner.public()),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(inner) => PublicKey::Secp256r1(*inner.public()),
            #[allow(unreachable_patterns)]
            _ => unreachable!(),
        }
    }

    /// Exports the secret scalar as bytes (for internal testing only).
    #[doc(hidden)]
    pub fn secret_scalar(&self) -> Vec<u8> {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(inner) => inner.secret().as_bytes().to_vec(),
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(inner) => inner.secret().to_bytes().as_slice().to_vec(),
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(inner) => inner.secret().to_bytes().as_slice().to_vec(),
            #[allow(unreachable_patterns)]
            _ => Vec::new(),
        }
    }

    /// Returns the curve identifier of this key.
    pub fn curve_id(&self) -> CurveId {
        match self {
            #[cfg(feature = "curve25519")]
            Self::Curve25519(_) => CurveId::Curve25519,
            #[cfg(feature = "secp256k1")]
            Self::Secp256k1(_) => CurveId::Secp256k1,
            #[cfg(feature = "secp256r1")]
            Self::Secp256r1(_) => CurveId::Secp256r1,
            #[allow(unreachable_patterns)]
            _ => unreachable!(),
        }
    }
}