//! Key exchange trait defining the ECMQV protocol interface.
//!
//! This module provides the `KeyExchange` trait, which abstracts the two
//! ECMQV modes: `EphemeralOnly` and `EphemeralStatic`. The trait is designed
//! to be implemented by those concrete types.

#![allow(unreachable_patterns)]
use alloc::vec::Vec;
use rand_core::{CryptoRng, RngCore};
use zeroize::Zeroizing;

use crate::error::AgreementError;
use crate::key::{EphemeralKey, PublicKey, StaticKey};
use crate::kdf::{derive_shared_secret, KdfParams, SharedInfo};
use crate::params::CurveParameters;

/// Core trait for ECMQV key exchange.
///
/// Implementations provide either ephemeral-only or ephemeral+static modes.
/// The shared secret is always a `Zeroizing<Vec<u8>>` containing the derived
/// symmetric key after applying the KDF.
///
/// # Examples
/// ```
/// # use ecmqv_rs::{EphemeralOnly, KeyExchange, CurveParameters, KdfParams, SharedInfo, HashFunction};
/// # use rand::rngs::OsRng;
/// let mut rng = OsRng;
/// let params = CurveParameters::Curve25519;
/// let (alice_eph, alice_pub) = EphemeralOnly::generate_ephemeral(&mut rng, &params).unwrap();
/// let (bob_eph, bob_pub) = EphemeralOnly::generate_ephemeral(&mut rng, &params).unwrap();
///
/// let alice = EphemeralOnly::new(alice_eph, params);
/// let bob = EphemeralOnly::new(bob_eph, params);
///
/// let info = SharedInfo {
///     id_a: b"alice",
///     id_b: b"bob",
///     other_data: b"session123",
/// };
/// let kdf = KdfParams {
///     hash_fn: HashFunction::Sha256,
///     output_len: 32,
///     label: None,
///     salt: None,
/// };
///
/// let alice_secret = alice.agree(&bob_pub, None, &info, kdf.clone()).unwrap();
/// let bob_secret = bob.agree(&alice_pub, None, &info, kdf).unwrap();
/// assert_eq!(alice_secret.as_slice(), bob_secret.as_slice());
/// ```
pub trait KeyExchange {
    /// The type of the shared secret produced by `agree`. Always a zeroized byte vector.
    type SharedSecret;

    /// Generates an ephemeral key pair for the given curve.
    ///
    /// # Parameters
    /// - `rng`: cryptographically secure random number generator
    /// - `params`: curve parameters (must be supported)
    ///
    /// # Returns
    /// A tuple `(ephemeral_key, public_key)` where the public key can be transmitted.
    /// The ephemeral key must be kept secret and destroyed after use.
    ///
    /// # Errors
    /// Returns `UnsupportedCurve` if the curve is not enabled.
    fn generate_ephemeral<R: RngCore + CryptoRng>(
        rng: &mut R,
        params: &CurveParameters,
    ) -> Result<(EphemeralKey, PublicKey), AgreementError>;

    /// Performs the ECMQV key agreement.
    ///
    /// # Parameters
    /// - `peer_ephemeral`: the peer’s ephemeral public key
    /// - `peer_static`: optional peer static public key (required if this side uses a static key)
    /// - `shared_info`: contextual information (identifiers, session data) for domain separation
    /// - `kdf_params`: parameters for the key derivation function
    ///
    /// # Returns
    /// The derived shared secret as a zeroized byte vector.
    ///
    /// # Errors
    /// - `ParameterMismatch` if keys belong to different curves or a required static key is missing.
    /// - `InvalidPublicKey` if the peer's public key fails validation.
    /// - `SharedPointIsIdentity` if the computed shared point is the point at infinity.
    /// - `KdfError` if the KDF fails.
    fn agree(
        &self,
        peer_ephemeral: &PublicKey,
        peer_static: Option<&PublicKey>,
        shared_info: &SharedInfo,
        kdf_params: KdfParams,
    ) -> Result<Self::SharedSecret, AgreementError>;
}

/// A zeroizable byte vector holding the derived shared secret.
///
/// This type is used as the output of `KeyExchange::agree`. It ensures
/// that the secret is automatically zeroized when dropped, reducing the
/// risk of accidental leakage.
pub type SharedSecret = Zeroizing<Vec<u8>>;

/// Helper function to verify that all keys belong to the same curve.
///
/// This is used internally to enforce curve consistency before performing
/// the agreement computation.
#[doc(hidden)]
pub fn check_curve_consistency(
    own_eph: &EphemeralKey,
    own_static: Option<&StaticKey>,
    peer_eph: &PublicKey,
    peer_static: Option<&PublicKey>,
    expected_curve: &CurveParameters,
) -> Result<(), AgreementError> {
    if own_eph.curve_id() != expected_curve.curve_id() {
        return Err(AgreementError::ParameterMismatch);
    }
    if let Some(s) = own_static {
        if s.curve_id() != expected_curve.curve_id() {
            return Err(AgreementError::ParameterMismatch);
        }
    }
    if peer_eph.curve_id() != expected_curve.curve_id() {
        return Err(AgreementError::ParameterMismatch);
    }
    if let Some(p) = peer_static {
        if p.curve_id() != expected_curve.curve_id() {
            return Err(AgreementError::ParameterMismatch);
        }
    }
    Ok(())
}

/// Helper function to perform the full agreement workflow.
///
/// It validates keys, computes the shared point, and applies the KDF.
/// This function is used internally by the public `agree` methods of
/// `EphemeralOnly` and `EphemeralStatic`.
#[doc(hidden)]
pub fn perform_agreement(
    own_eph: &EphemeralKey,
    own_static: Option<&StaticKey>,
    peer_eph: &PublicKey,
    peer_static: Option<&PublicKey>,
    curve: &CurveParameters,
    shared_info: &SharedInfo,
    kdf_params: KdfParams,
) -> Result<SharedSecret, AgreementError> {
    // Validate peer public keys (security critical)
    peer_eph.validate()?;
    if let Some(peer_static) = peer_static {
        peer_static.validate()?;
    }

    // Compute the raw shared point
    let shared_point_bytes = crate::agreement::compute_shared_point(
        own_eph,
        own_static,
        peer_eph,
        peer_static,
        curve,
    )?;

    // Derive the final secret using the KDF
    let secret = derive_shared_secret(&shared_point_bytes, shared_info, &kdf_params)?;
    Ok(secret)
}