//! Key derivation function (HKDF) for ECMQV shared secrets.

#![allow(unreachable_patterns)]

use alloc::vec;
use alloc::vec::Vec;
use zeroize::Zeroizing;

use crate::error::AgreementError;

/// Supported hash functions for HKDF.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum HashFunction {
    /// SHA-256 with 32-byte output.
    Sha256,
    /// SHA-512 with 64-byte output.
    Sha512,
    /// BLAKE3 with 32-byte output by default, but any length up to 2^64-1 is supported.
    Blake3,
}

impl HashFunction {
    /// Returns the output length in bytes of the underlying hash function.
    pub fn hash_len(&self) -> usize {
        match self {
            Self::Sha256 => 32,
            Self::Sha512 => 64,
            Self::Blake3 => 32,
        }
    }
}

/// Parameters for the key derivation function (HKDF).
#[derive(Debug, Clone)]
pub struct KdfParams {
    /// The hash function to use.
    pub hash_fn: HashFunction,
    /// Desired output length in bytes (must be <= 255 * hash_len).
    pub output_len: usize,
    /// Optional application-specific label (included in the info string).
    pub label: Option<&'static [u8]>,
    /// Optional salt. If `None`, a zero-filled salt of length `hash_len` is used.
    /// For most ECMQV uses, salt should be `None`.
    pub salt: Option<&'static [u8]>,
}

/// Context information that MUST be bound into the KDF to prevent
/// unknown key-share attacks and ensure domain separation.
pub struct SharedInfo<'a> {
    /// Identifier of the first party (e.g., certificate hash, DER‑encoded name).
    pub id_a: &'a [u8],
    /// Identifier of the second party.
    pub id_b: &'a [u8],
    /// Additional contextual data (e.g., session ID, protocol version).
    pub other_data: &'a [u8],
}

impl<'a> SharedInfo<'a> {
    /// Serializes the shared info into a single byte string suitable for HKDF `info`.
    ///
    /// Format:
    /// `"ecmqv-kdf-v1" || len(id_a) as u16 BE || id_a || len(id_b) as u16 BE || id_b || other_data || label`
    ///
    /// This ensures domain separation and prevents cross-protocol attacks.
    fn to_info_bytes(&self, label: Option<&[u8]>) -> Vec<u8> {
        const DOMAIN_SEP: &[u8] = b"ecmqv-kdf-v1";
        let mut buf = Vec::with_capacity(
            DOMAIN_SEP.len()
                + 2 + self.id_a.len()
                + 2 + self.id_b.len()
                + self.other_data.len()
                + label.map(|l| l.len()).unwrap_or(0),
        );
        buf.extend_from_slice(DOMAIN_SEP);
        buf.extend_from_slice(&(self.id_a.len() as u16).to_be_bytes());
        buf.extend_from_slice(self.id_a);
        buf.extend_from_slice(&(self.id_b.len() as u16).to_be_bytes());
        buf.extend_from_slice(self.id_b);
        buf.extend_from_slice(self.other_data);
        if let Some(label) = label {
            buf.extend_from_slice(label);
        }
        buf
    }
}

/// Derives a shared secret from the raw ECMQV shared point and the context.
///
/// This function applies HKDF using the chosen hash function, salt, and info
/// derived from the `SharedInfo` structure. The output is zeroized on drop.
///
/// # Parameters
/// - `shared_secret`: the raw shared point bytes from ECMQV.
/// - `info`: contextual identifiers for domain separation.
/// - `params`: KDF parameters (hash, output length, label, salt).
///
/// # Returns
/// A `Zeroizing<Vec<u8>>` containing the derived key material.
///
/// # Errors
/// - `InvalidKeyLength` if `output_len` is 0 or greater than `255 * hash_len`.
/// - `UnsupportedHashFunction` if the chosen hash is not enabled.
/// - `KdfError` if HKDF expansion fails.
pub fn derive_shared_secret(
    shared_secret: &[u8],
    info: &SharedInfo,
    params: &KdfParams,
) -> Result<Zeroizing<Vec<u8>>, AgreementError> {
    let hash_len = params.hash_fn.hash_len();

    if params.output_len == 0 || params.output_len > 255 * hash_len {
        return Err(AgreementError::InvalidKeyLength);
    }

    let info_bytes = info.to_info_bytes(params.label);

    match params.hash_fn {
        #[cfg(feature = "sha2")]
        HashFunction::Sha256 => {
            use sha2::Sha256;
            let salt = params.salt.unwrap_or(&[0u8; 32]);
            let hkdf = hkdf::Hkdf::<Sha256>::new(Some(salt), shared_secret);
            let mut okm = vec![0u8; params.output_len];
            hkdf.expand(&info_bytes, &mut okm)
                .map_err(|_| AgreementError::KdfError)?;
            Ok(Zeroizing::new(okm))
        }
        #[cfg(feature = "sha2")]
        HashFunction::Sha512 => {
            use sha2::Sha512;
            let salt = params.salt.unwrap_or(&[0u8; 64]);
            let hkdf = hkdf::Hkdf::<Sha512>::new(Some(salt), shared_secret);
            let mut okm = vec![0u8; params.output_len];
            hkdf.expand(&info_bytes, &mut okm)
                .map_err(|_| AgreementError::KdfError)?;
            Ok(Zeroizing::new(okm))
        }
        #[cfg(feature = "blake3")]
        HashFunction::Blake3 => {
            use blake3::Hasher;
            let salt = params.salt.unwrap_or(&[0u8; 32]);
            if salt.len() != 32 {
                return Err(AgreementError::InvalidKeyLength);
            }
            let salt_array: &[u8; 32] = salt
                .try_into()
                .map_err(|_| AgreementError::InvalidKeyLength)?;
            let prk = blake3::keyed_hash(salt_array, shared_secret);
            let prk_bytes = prk.as_bytes();
            let mut okm = Vec::with_capacity(params.output_len);
            let mut counter: u8 = 1;
            let mut t_prev = [0u8; 32];
            while okm.len() < params.output_len {
                let mut hasher = Hasher::new_keyed(prk_bytes);
                if counter > 1 {
                    hasher.update(&t_prev);
                }
                hasher.update(&info_bytes);
                hasher.update(&[counter]);
                let t = hasher.finalize();
                let needed = params.output_len - okm.len();
                let take = needed.min(32);
                okm.extend_from_slice(&t.as_bytes()[..take]);
                t_prev.copy_from_slice(t.as_bytes());
                counter = counter.checked_add(1).ok_or(AgreementError::KdfError)?;
            }
            Ok(Zeroizing::new(okm))
        }
        _ => Err(AgreementError::UnsupportedHashFunction),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_hkdf_sha256() {
        let shared = b"test shared secret for HKDF";
        let info = SharedInfo {
            id_a: b"alice",
            id_b: b"bob",
            other_data: b"session123",
        };
        let params = KdfParams {
            hash_fn: HashFunction::Sha256,
            output_len: 32,
            label: Some(b"test"),
            salt: None,
        };
        let result = derive_shared_secret(shared, &info, &params).unwrap();
        assert_eq!(result.len(), 32);
    }
}