//! Asymmetric key handling and wrappers

use ::num_traits::FromPrimitive;

use super::Error;
use crate::{
    config::Config,
    enc::{Random, Secret},
};

/// Public key ID
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct KeyID(pub u16);

impl KeyID {
    /// Length of the Key ID in bytes
    pub const fn len() -> usize {
        2
    }
    /// Maximum possible KeyID
    pub const MAX: u16 = u16::MAX;
    /// Serialize into raw bytes
    pub fn serialize(&self, out: &mut [u8; KeyID::len()]) {
        out.copy_from_slice(&self.0.to_le_bytes());
    }
}

impl TryFrom<&str> for KeyID {
    type Error = ::std::io::Error;
    fn try_from(raw: &str) -> Result<Self, Self::Error> {
        if let Ok(id_u16) = raw.parse::<u16>() {
            return Ok(KeyID(id_u16));
        }
        return Err(::std::io::Error::new(
            ::std::io::ErrorKind::InvalidData,
            "KeyID must be between 0 and 65535",
        ));
    }
}
impl ::std::fmt::Display for KeyID {
    fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Capabilities of each key
#[derive(Debug, Clone, Copy)]
pub enum KeyCapabilities {
    /// signing *only*
    Sign,
    /// encrypt *only*
    Encrypt,
    /// key exchange *only*
    Exchange,
    /// both sign and encrypt
    SignEncrypt,
    /// both signing and key exchange
    SignExchange,
    /// both encrypt and key exchange
    EncryptExchange,
    /// All: sign, encrypt, Key Exchange
    SignEncryptExchage,
}
impl KeyCapabilities {
    /// Check if this key supports eky exchage
    pub fn has_exchange(&self) -> bool {
        match self {
            KeyCapabilities::Exchange
            | KeyCapabilities::SignExchange
            | KeyCapabilities::SignEncryptExchage => true,
            _ => false,
        }
    }
}

/// Kind of key used in the handshake
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    ::num_derive::FromPrimitive,
    ::strum_macros::EnumString,
    ::strum_macros::IntoStaticStr,
)]
#[non_exhaustive]
#[repr(u8)]
pub enum KeyKind {
    /// Ed25519 Public key (sign only)
    #[strum(serialize = "ed25519")]
    Ed25519 = 0,
    /// X25519 Public key (key exchange)
    #[strum(serialize = "x25519")]
    X25519,
}
// FIXME: actually check this
const MIN_KEY_SIZE: usize = 32;
impl KeyKind {
    /// return the expected length of the public key
    pub fn pub_len(&self) -> usize {
        match self {
            // FIXME: 99% wrong size
            KeyKind::Ed25519 => ::ring::signature::ED25519_PUBLIC_KEY_LEN,
            KeyKind::X25519 => 32,
        }
    }
    /// Get the capabilities of this key type
    pub fn capabilities(&self) -> KeyCapabilities {
        match self {
            KeyKind::Ed25519 => KeyCapabilities::Sign,
            KeyKind::X25519 => KeyCapabilities::Exchange,
        }
    }
    /// Returns the key exchanges supported by this key
    pub fn key_exchanges(&self) -> &'static [KeyExchangeKind] {
        const EMPTY: [KeyExchangeKind; 0] = [];
        const X25519_KEY_EXCHANGES: [KeyExchangeKind; 1] =
            [KeyExchangeKind::X25519DiffieHellman];
        match self {
            KeyKind::Ed25519 => &EMPTY,
            KeyKind::X25519 => &X25519_KEY_EXCHANGES,
        }
    }
    /// generate new keypair
    pub fn new_keypair(
        &self,
        rnd: &Random,
    ) -> Result<(PrivKey, PubKey), Error> {
        PubKey::new_keypair(*self, rnd)
    }
}

/// Kind of key exchange
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    ::num_derive::FromPrimitive,
    ::strum_macros::EnumString,
    ::strum_macros::IntoStaticStr,
)]
#[non_exhaustive]
#[repr(u8)]
pub enum KeyExchangeKind {
    /// X25519 Public key
    #[strum(serialize = "x25519diffiehellman")]
    X25519DiffieHellman = 0,
}
impl KeyExchangeKind {
    /// The serialize length of the field
    pub fn len() -> usize {
        1
    }
    /// Build a new keypair for key exchange
    pub fn new_keypair(
        &self,
        rnd: &Random,
    ) -> Result<(ExchangePrivKey, ExchangePubKey), Error> {
        match self {
            KeyExchangeKind::X25519DiffieHellman => {
                let raw_priv = ::x25519_dalek::StaticSecret::new(rnd);
                let pub_key = ExchangePubKey::X25519(
                    ::x25519_dalek::PublicKey::from(&raw_priv),
                );
                let priv_key = ExchangePrivKey::X25519(raw_priv);
                Ok((priv_key, pub_key))
            }
            _ => Err(Error::UnsupportedKeyExchange),
        }
    }
}

/// Kind of public key in the handshake
#[derive(Debug, Copy, Clone, Eq, Hash, PartialEq)]
#[allow(missing_debug_implementations)]
#[non_exhaustive]
pub enum PubKey {
    /// Keys to be used only in key exchanges, not for signing
    Exchange(ExchangePubKey),
    /// Keys to be used only for signing
    Signing,
}

impl PubKey {
    /// Get the serialized key length
    pub fn len(&self) -> usize {
        match self {
            PubKey::Exchange(ex) => ex.len(),
            PubKey::Signing => todo!(),
        }
    }
    /// return the kind of public key
    pub fn kind(&self) -> KeyKind {
        match self {
            // FIXME: lie, we don't fully support this
            PubKey::Signing => KeyKind::Ed25519,
            PubKey::Exchange(ex) => ex.kind(),
        }
    }
    /// generate new keypair
    fn new_keypair(
        kind: KeyKind,
        rnd: &Random,
    ) -> Result<(PrivKey, PubKey), Error> {
        match kind {
            KeyKind::Ed25519 => todo!(),
            KeyKind::X25519 => {
                let (priv_key, pub_key) =
                    KeyExchangeKind::X25519DiffieHellman.new_keypair(rnd)?;
                Ok((PrivKey::Exchange(priv_key), PubKey::Exchange(pub_key)))
            }
        }
    }
    /// serialize the key into the buffer
    /// NOTE: Assumes there is enough space
    pub fn serialize_into(&self, out: &mut [u8]) {
        assert!(
            out.len() >= 1 + self.kind().pub_len(),
            "Not enough out buffer",
        );
        out[0] = self.kind() as u8;
        match self {
            PubKey::Signing => {
                ::tracing::error!("serializing ed25519 not supported");
                return;
            }
            PubKey::Exchange(ex) => ex.serialize_into(&mut out[1..]),
        }
    }
    /// Try to deserialize the pubkey from raw bytes
    /// on success returns the public key and the number of parsed bytes
    pub fn deserialize(raw: &[u8]) -> Result<(Self, usize), Error> {
        if raw.len() < 1 + MIN_KEY_SIZE {
            return Err(Error::NotEnoughData(0));
        }
        let kind: KeyKind = match KeyKind::from_u8(raw[0]) {
            Some(kind) => kind,
            None => return Err(Error::UnsupportedKey(1)),
        };
        if raw.len() < 1 + kind.pub_len() {
            return Err(Error::NotEnoughData(1));
        }
        match kind {
            KeyKind::Ed25519 => {
                ::tracing::error!("ed25519 keys are not yet supported");
                return Err(Error::Parsing);
            }
            KeyKind::X25519 => {
                let pub_key: ::x25519_dalek::PublicKey =
                    match ::bincode::deserialize(&raw[1..(1 + kind.pub_len())])
                    {
                        Ok(pub_key) => pub_key,
                        Err(_) => return Err(Error::Parsing),
                    };
                Ok((
                    PubKey::Exchange(ExchangePubKey::X25519(pub_key)),
                    kind.pub_len(),
                ))
            }
        }
    }
}

/// Kind of private key in the handshake
#[derive(Clone)]
#[allow(missing_debug_implementations)]
#[non_exhaustive]
pub enum PrivKey {
    /// Keys to be used only in key exchanges, not for signing
    Exchange(ExchangePrivKey),
    /// Keys to be used only for signing
    // TODO: implement ed25519
    Signing,
}

impl PrivKey {
    /// Get the serialized key length
    pub fn len(&self) -> usize {
        match self {
            PrivKey::Exchange(ex) => ex.len(),
            PrivKey::Signing => todo!(),
        }
    }
    /// serialize the key into the buffer
    /// NOTE: Assumes there is enough space
    pub fn serialize_into(&self, out: &mut [u8]) {
        match self {
            PrivKey::Exchange(ex) => ex.serialize_into(out),
            PrivKey::Signing => todo!(),
        }
    }
}

/// Ephemeral private keys
#[derive(Clone)]
#[allow(missing_debug_implementations)]
#[non_exhaustive]
pub enum ExchangePrivKey {
    /// X25519(Curve25519) used for key exchange
    X25519(::x25519_dalek::StaticSecret),
}

impl ExchangePrivKey {
    /// Get the serialized key length
    pub fn len(&self) -> usize {
        match self {
            ExchangePrivKey::X25519(_) => KeyKind::X25519.pub_len(),
        }
    }
    /// Get the kind of key
    pub fn kind(&self) -> KeyKind {
        match self {
            ExchangePrivKey::X25519(_) => KeyKind::X25519,
        }
    }
    /// Run the key exchange between two keys of the same kind
    pub fn key_exchange(
        &self,
        exchange: KeyExchangeKind,
        pub_key: ExchangePubKey,
    ) -> Result<Secret, Error> {
        match self {
            ExchangePrivKey::X25519(priv_key) => {
                if exchange != KeyExchangeKind::X25519DiffieHellman {
                    return Err(Error::UnsupportedKeyExchange);
                }
                let ExchangePubKey::X25519(inner_pub_key) = pub_key;
                let shared_secret = priv_key.diffie_hellman(&inner_pub_key);
                Ok(shared_secret.into())
            }
        }
    }
    /// serialize the key into the buffer
    /// NOTE: Assumes there is enough space
    pub fn serialize_into(&self, out: &mut [u8]) {
        match self {
            ExchangePrivKey::X25519(key) => {
                out[..32].copy_from_slice(&key.to_bytes());
            }
        }
    }
}

/// all Ephemeral Public keys
#[derive(Debug, Copy, Clone, Eq, Hash, PartialEq)]
#[non_exhaustive]
pub enum ExchangePubKey {
    /// X25519(Curve25519) used for key exchange
    X25519(::x25519_dalek::PublicKey),
}

impl ExchangePubKey {
    /// Get the serialized key length
    pub fn len(&self) -> usize {
        match self {
            ExchangePubKey::X25519(_) => KeyKind::X25519.pub_len(),
        }
    }
    /// Get the kind of key
    pub fn kind(&self) -> KeyKind {
        match self {
            ExchangePubKey::X25519(_) => KeyKind::X25519,
        }
    }
    /// serialize the key into the buffer
    /// NOTE: Assumes there is enough space
    pub fn serialize_into(&self, out: &mut [u8]) {
        match self {
            ExchangePubKey::X25519(pk) => {
                let bytes = pk.as_bytes();
                assert!(bytes.len() == 32, "x25519 should have been 32 bytes");
                out[..32].copy_from_slice(bytes);
            }
        }
    }
    /// Load public key used for key exchange from it raw bytes
    /// The riesult is "unparsed" since we don't verify
    /// the actual key
    pub fn from_slice(raw: &[u8]) -> Result<(Self, usize), Error> {
        if raw.len() < 1 + MIN_KEY_SIZE {
            return Err(Error::NotEnoughData(0));
        }
        match KeyKind::from_u8(raw[0]) {
            Some(kind) => match kind {
                KeyKind::Ed25519 => {
                    ::tracing::error!("ed25519 keys are not yet supported");
                    return Err(Error::Parsing);
                }
                KeyKind::X25519 => {
                    let pub_key: ::x25519_dalek::PublicKey =
                        match ::bincode::deserialize(
                            &raw[1..(1 + kind.pub_len())],
                        ) {
                            Ok(pub_key) => pub_key,
                            Err(_) => return Err(Error::Parsing),
                        };
                    Ok((ExchangePubKey::X25519(pub_key), kind.pub_len()))
                }
            },
            None => {
                return Err(Error::Parsing);
            }
        }
    }
}

/// Select the best key exchange from our supported list
/// and the other endpoint supported list.
/// Give priority to our list
pub fn server_select_key_exchange(
    cfg: &Config,
    client_supported: &Vec<KeyExchangeKind>,
) -> Option<KeyExchangeKind> {
    cfg.key_exchanges
        .iter()
        .find(|k| client_supported.contains(k))
        .copied()
}
/// Select the best key exchange 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_key_exchange(
    cfg: &Config,
    server_supported: &Vec<KeyExchangeKind>,
) -> Option<KeyExchangeKind> {
    server_supported
        .iter()
        .find(|k| cfg.key_exchanges.contains(k))
        .copied()
}