//! Symmetric cypher stuff

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

/// List of possible Ciphers
#[derive(
    Debug,
    Copy,
    Clone,
    PartialEq,
    ::num_derive::FromPrimitive,
    ::strum_macros::EnumString,
    ::strum_macros::IntoStaticStr,
)]
#[repr(u8)]
pub enum Kind {
    /// XChaCha20_Poly1305
    #[strum(serialize = "xchacha20poly1305")]
    XChaCha20Poly1305 = 0,
}

impl Kind {
    /// length of the serialized id for the cipher kind field
    pub const fn len() -> usize {
        1
    }
    /// required length of the nonce
    pub fn nonce_len(&self) -> NonceLen {
        Nonce::len()
    }
    /// required length of the key
    pub fn key_len(&self) -> usize {
        use ::chacha20poly1305::{KeySizeUser, XChaCha20Poly1305};
        XChaCha20Poly1305::key_size()
    }
    /// Length of the authentication tag
    pub fn tag_len(&self) -> TagLen {
        // TODO: how the hell do I take this from ::chacha20poly1305?
        TagLen(::ring::aead::CHACHA20_POLY1305.tag_len())
    }
}

/// Additional Authenticated Data
#[derive(Debug)]
pub struct AAD<'a>(pub &'a [u8]);

/// strong typedef for header length
/// aka: nonce length in the encrypted data)
#[derive(Debug, Copy, Clone)]
pub struct NonceLen(pub usize);
/// strong typedef for the Tag length
/// aka: cryptographic authentication tag length at the end
/// of the encrypted data
#[derive(Debug, Copy, Clone)]
pub struct TagLen(pub usize);

/// actual ciphers
enum Cipher {
    /// Cipher XChaha20_Poly1305
    XChaCha20Poly1305(XChaCha20Poly1305),
}

impl Cipher {
    /// Build a new Cipher
    fn new(kind: Kind, secret: Secret) -> Self {
        match kind {
            Kind::XChaCha20Poly1305 => {
                Self::XChaCha20Poly1305(XChaCha20Poly1305::new(secret))
            }
        }
    }
    pub fn kind(&self) -> Kind {
        match self {
            Cipher::XChaCha20Poly1305(_) => Kind::XChaCha20Poly1305,
        }
    }
    fn nonce_len(&self) -> NonceLen {
        match self {
            Cipher::XChaCha20Poly1305(_) => Nonce::len(),
        }
    }
    fn tag_len(&self) -> TagLen {
        match self {
            Cipher::XChaCha20Poly1305(_) => {
                // TODO: how the hell do I take this from ::chacha20poly1305?
                TagLen(::ring::aead::CHACHA20_POLY1305.tag_len())
            }
        }
    }
    fn decrypt<'a>(
        &self,
        aad: AAD,
        raw_data: &'a mut [u8],
    ) -> Result<&'a [u8], Error> {
        match self {
            Cipher::XChaCha20Poly1305(cipher) => {
                use ::chacha20poly1305::{
                    aead::generic_array::GenericArray, AeadInPlace,
                };
                let final_len: usize = {
                    if raw_data.len() <= self.overhead() {
                        return Err(Error::NotEnoughData(raw_data.len()));
                    }
                    let (nonce_bytes, data_and_tag) =
                        raw_data.split_at_mut(Nonce::len().0);
                    let (data_notag, tag_bytes) = data_and_tag.split_at_mut(
                        data_and_tag.len()
                            - ::ring::aead::CHACHA20_POLY1305.tag_len(),
                    );
                    let nonce = GenericArray::from_slice(nonce_bytes);
                    let tag = GenericArray::from_slice(tag_bytes);
                    let maybe = cipher.cipher.decrypt_in_place_detached(
                        nonce.into(),
                        aad.0,
                        data_notag,
                        tag,
                    );
                    if maybe.is_err() {
                        return Err(Error::Decrypt);
                    }
                    data_notag.len()
                };
                //data.drain(..Nonce::len());
                //data.truncate(final_len);
                Ok(&raw_data[Nonce::len().0..Nonce::len().0 + final_len])
            }
        }
    }
    fn overhead(&self) -> usize {
        match self {
            Cipher::XChaCha20Poly1305(_) => {
                let cipher = Kind::XChaCha20Poly1305;
                cipher.nonce_len().0 + cipher.tag_len().0
            }
        }
    }
    fn encrypt(
        &mut self,
        nonce: &Nonce,
        aad: AAD,
        data: &mut [u8],
    ) -> Result<(), Error> {
        // FIXME: check minimum buffer size
        match self {
            Cipher::XChaCha20Poly1305(cipher) => {
                use ::chacha20poly1305::AeadInPlace;
                let tag_len: usize = ::ring::aead::CHACHA20_POLY1305.tag_len();
                let data_len_notag = data.len() - tag_len;
                // write nonce
                data[..Nonce::len().0].copy_from_slice(nonce.as_bytes());

                // encrypt data
                match cipher.cipher.encrypt_in_place_detached(
                    nonce.as_bytes().into(),
                    aad.0,
                    &mut data[Nonce::len().0..data_len_notag],
                ) {
                    Ok(tag) => {
                        data[data_len_notag..].copy_from_slice(tag.as_slice());
                        Ok(())
                    }
                    Err(_) => Err(Error::Encrypt),
                }
            }
        }
    }
}

/// Receive only cipher
pub struct CipherRecv(Cipher);
impl ::core::fmt::Debug for CipherRecv {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden cipher recv]", f)
    }
}

impl CipherRecv {
    /// Build a new Cipher
    pub fn new(kind: Kind, secret: Secret) -> Self {
        Self(Cipher::new(kind, secret))
    }
    /// Get the length of the nonce  for this cipher
    pub fn nonce_len(&self) -> NonceLen {
        self.0.nonce_len()
    }
    /// Get the length of the nonce  for this cipher
    pub fn tag_len(&self) -> TagLen {
        self.0.tag_len()
    }
    /// Decrypt a paket. Nonce and Tag are taken from the packet,
    /// while you need to provide AAD (Additional Authenticated Data)
    pub fn decrypt<'a>(
        &self,
        aad: AAD,
        data: &'a mut [u8],
    ) -> Result<&'a [u8], Error> {
        self.0.decrypt(aad, data)
    }
    /// return the underlying cipher id
    pub fn kind(&self) -> Kind {
        self.0.kind()
    }
}

/// Send only cipher
pub struct CipherSend {
    nonce: Nonce,
    cipher: Cipher,
}
impl ::core::fmt::Debug for CipherSend {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden cipher send]", f)
    }
}

impl CipherSend {
    /// Build a new Cipher
    pub fn new(kind: Kind, secret: Secret, rand: &Random) -> Self {
        Self {
            nonce: Nonce::new(rand),
            cipher: Cipher::new(kind, secret),
        }
    }
    /// Encrypt the given data
    pub fn encrypt(&mut self, aad: AAD, data: &mut [u8]) -> Result<(), Error> {
        let old_nonce = self.nonce.advance();
        self.cipher.encrypt(&old_nonce, aad, data)?;
        Ok(())
    }
    /// return the underlying cipher id
    pub fn kind(&self) -> Kind {
        self.cipher.kind()
    }
}

/// XChaCha20Poly1305 cipher
struct XChaCha20Poly1305 {
    cipher: ::chacha20poly1305::XChaCha20Poly1305,
}

impl XChaCha20Poly1305 {
    /// Initialize a new ChaChaPoly20 cipher with a random nonce
    fn new(key: Secret) -> Self {
        use ::chacha20poly1305::{KeyInit, XChaCha20Poly1305};
        Self {
            cipher: XChaCha20Poly1305::new(key.as_ref().into()),
        }
    }
}

//
// TODO: Merge crate::{enc::sym::Nonce, connection::handshake::dirsync::Nonce}
//

#[derive(Debug, Copy, Clone)]
#[repr(C)]
#[allow(missing_debug_implementations)]
struct NonceNum {
    high: u32,
    low: u64,
}
/// Nonce with sequence for chach20_apoly1305
#[derive(Copy, Clone)]
#[repr(C)]
pub union Nonce {
    num: NonceNum,
    raw: [u8; Self::len().0],
}

impl ::core::fmt::Debug for Nonce {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        #[allow(unsafe_code)]
        unsafe {
            ::core::fmt::Debug::fmt(&self.num, f)
        }
    }
}

impl Nonce {
    /// Generate a new random Nonce
    pub fn new(rand: &Random) -> Self {
        let mut raw = [0; Self::len().0];
        rand.fill(&mut raw);
        Self { raw }
    }
    /// Length of this nonce in bytes
    pub const fn len() -> NonceLen {
        // FIXME: was:12. xchacha20poly1305 requires 24.
        // but we should change keys much earlier than that, and our
        // nonces are not random, but sequential.
        // we should change keys every 2^30 bytes to be sure (stream max window)
        return NonceLen(24);
    }
    /// Get reference to the nonce bytes
    pub fn as_bytes(&self) -> &[u8] {
        #[allow(unsafe_code)]
        unsafe {
            &self.raw
        }
    }
    /// Create Nonce from array
    pub fn from_slice(raw: [u8; Self::len().0]) -> Self {
        Self { raw }
    }
    /// Go to the next nonce
    pub fn advance(&mut self) -> Self {
        let old_nonce = self.clone();
        #[allow(unsafe_code)]
        unsafe {
            let old_low = self.num.low;
            self.num.low = self.num.low + 1;
            if self.num.low < old_low {
                self.num.high = self.num.high;
            }
        }
        old_nonce
    }
}

/// Select the best cipher from our supported list
/// and the other endpoint supported list.
/// Give priority to our list
pub fn server_select_cipher(
    cfg: &Config,
    client_supported: &Vec<Kind>,
) -> Option<Kind> {
    cfg.ciphers
        .iter()
        .find(|c| client_supported.contains(c))
        .copied()
}
/// Select the best cipher 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_cipher(
    cfg: &Config,
    server_supported: &Vec<Kind>,
) -> Option<Kind> {
    server_supported
        .iter()
        .find(|c| cfg.ciphers.contains(c))
        .copied()
}