//! Hash-based Key Derivation Function
//! We just repackage other crates

use ::sha3::Sha3_256;
use ::zeroize::Zeroize;

use crate::{config::Config, enc::sym::Secret};

/// Kind of HKDF
#[derive(Debug, Copy, Clone, PartialEq, ::num_derive::FromPrimitive)]
#[non_exhaustive]
#[repr(u8)]
pub enum HkdfKind {
    /// Sha3
    Sha3 = 0,
}
impl HkdfKind {
    /// Length of the serialized type
    pub const fn len() -> usize {
        1
    }
}

/// Generic wrapper on Hkdfs
#[derive(Clone)]
pub enum Hkdf {
    /// Sha3 based
    Sha3(HkdfSha3),
}

// Fake debug implementation to avoid leaking secrets
impl ::core::fmt::Debug for Hkdf {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden hkdf]", f)
    }
}

impl Hkdf {
    /// New Hkdf
    pub fn new(kind: HkdfKind, salt: &[u8], key: Secret) -> Self {
        match kind {
            HkdfKind::Sha3 => Self::Sha3(HkdfSha3::new(salt, key)),
        }
    }
    /// Get a secret generated from the key and a given context
    pub fn get_secret(&self, context: &[u8]) -> Secret {
        match self {
            Hkdf::Sha3(sha3) => sha3.get_secret(context),
        }
    }
}

// Hack & tricks:
// HKDF are pretty important, but this lib don't zero out the data.
// we can't use #[derive(Zeroing)] either.
// So we craete a union with a Zeroing object, and drop both manually.

// TODO: move this to Hkdf instead of Sha3

#[derive(Zeroize)]
#[zeroize(drop)]
struct Zeroable([u8; ::core::mem::size_of::<::hkdf::Hkdf<Sha3_256>>()]);

union HkdfInner {
    hkdf: ::core::mem::ManuallyDrop<::hkdf::Hkdf<Sha3_256>>,
    zeroable: ::core::mem::ManuallyDrop<Zeroable>,
}

impl Drop for HkdfInner {
    fn drop(&mut self) {
        #[allow(unsafe_code)]
        unsafe {
            drop(&mut self.hkdf);
            drop(&mut self.zeroable);
        }
    }
}
impl Clone for HkdfInner {
    fn clone(&self) -> Self {
        #[allow(unsafe_code)]
        unsafe {
            Self {
                hkdf: self.hkdf.clone(),
            }
        }
    }
}

/// Sha3 based HKDF
#[derive(Clone)]
pub struct HkdfSha3 {
    inner: HkdfInner,
}

impl HkdfSha3 {
    /// Instantiate a new HKDF with Sha3-256
    pub(crate) fn new(salt: &[u8], key: Secret) -> Self {
        let hkdf = ::hkdf::Hkdf::<Sha3_256>::new(Some(salt), key.as_ref());
        Self {
            inner: HkdfInner {
                hkdf: ::core::mem::ManuallyDrop::new(hkdf),
            },
        }
    }
    /// Get a secret generated from the key and a given context
    pub(crate) fn get_secret(&self, context: &[u8]) -> Secret {
        let mut out: [u8; 32] = [0; 32];
        #[allow(unsafe_code)]
        unsafe {
            self.inner.hkdf.expand(context, &mut out);
        }
        out.into()
    }
}

// Fake debug implementation to avoid leaking secrets
impl ::core::fmt::Debug for HkdfSha3 {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden hkdf]", f)
    }
}

/// Select the best hkdf from our supported list
/// and the other endpoint supported list.
/// Give priority to our list
pub fn server_select_hkdf(
    cfg: &Config,
    client_supported: &Vec<HkdfKind>,
) -> Option<HkdfKind> {
    cfg.hkdfs
        .iter()
        .find(|h| client_supported.contains(h))
        .copied()
}
/// Select the best hkdf from our supported list
/// and the other endpoint supported list.
/// Give priority to the server list
/// this is used only in the directory synchronized handshake
pub fn client_select_hkdf(
    cfg: &Config,
    server_supported: &Vec<HkdfKind>,
) -> Option<HkdfKind> {
    server_supported
        .iter()
        .find(|h| cfg.hkdfs.contains(h))
        .copied()
}