//! Symmetric cypher stuff

use super::Error;
use ::std::collections::VecDeque;
use ::zeroize::Zeroize;

/// Secret, used for keys.
/// Grants that on drop() we will zero out memory
#[derive(Zeroize, Clone)]
#[zeroize(drop)]
pub struct Secret([u8; 32]);
// Fake debug implementation to avoid leaking secrets
impl ::core::fmt::Debug for Secret {
    fn fmt(
        &self,
        f: &mut core::fmt::Formatter<'_>,
    ) -> Result<(), ::std::fmt::Error> {
        ::core::fmt::Debug::fmt("[hidden secret]", f)
    }
}

impl Secret {
    /// New randomly generated secret
    pub fn new_rand(rand: &::ring::rand::SystemRandom) -> Self {
        use ::ring::rand::SecureRandom;
        let mut ret = Self([0; 32]);
        rand.fill(&mut ret.0);
        ret
    }
    /// return a reference to the secret
    pub fn as_ref(&self) -> &[u8; 32] {
        &self.0
    }
}

impl From<[u8; 32]> for Secret {
    fn from(shared_secret: [u8; 32]) -> Self {
        Self(shared_secret)
    }
}

impl From<::x25519_dalek::SharedSecret> for Secret {
    fn from(shared_secret: ::x25519_dalek::SharedSecret) -> Self {
        Self(shared_secret.to_bytes())
    }
}

/// List of possible Ciphers
#[derive(Debug, Copy, Clone, PartialEq, ::num_derive::FromPrimitive)]
#[repr(u8)]
pub enum CipherKind {
    /// XChaCha20_Poly1305
    XChaCha20Poly1305 = 0,
}

impl CipherKind {
    /// required length of the nonce
    pub fn nonce_len(&self) -> usize {
        // TODO: how the hell do I take this from ::chacha20poly1305?
        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) -> usize {
        // TODO: how the hell do I take this from ::chacha20poly1305?
        ::ring::aead::CHACHA20_POLY1305.tag_len()
    }
}

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

/// Cipher direction, to make sure we don't reuse the same cipher
/// for both decrypting and encrypting
#[derive(Debug, Copy, Clone)]
#[repr(u8)]
pub enum CipherDirection {
    /// Receive, to decrypt only
    Recv = 0,
    /// Send, to encrypt only
    Send,
}

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

impl Cipher {
    /// Build a new Cipher
    fn new(kind: CipherKind, secret: Secret) -> Self {
        match kind {
            CipherKind::XChaCha20Poly1305 => {
                Self::XChaCha20Poly1305(XChaCha20Poly1305::new(secret))
            }
        }
    }
    fn nonce_len(&self) -> usize {
        match self {
            Cipher::XChaCha20Poly1305(_) => {
                // TODO: how the hell do I take this from ::chacha20poly1305?
                ::ring::aead::CHACHA20_POLY1305.nonce_len()
            }
        }
    }
    fn tag_len(&self) -> usize {
        match self {
            Cipher::XChaCha20Poly1305(_) => {
                // TODO: how the hell do I take this from ::chacha20poly1305?
                ::ring::aead::CHACHA20_POLY1305.tag_len()
            }
        }
    }
    fn decrypt(&self, aad: AAD, data: &mut VecDeque<u8>) -> Result<(), Error> {
        match self {
            Cipher::XChaCha20Poly1305(cipher) => {
                use ::chacha20poly1305::{
                    aead::generic_array::GenericArray, AeadInPlace,
                };
                let final_len: usize;
                {
                    let raw_data = data.as_mut_slices().0;
                    // FIXME: check min data length
                    let (nonce_bytes, data_and_tag) = raw_data.split_at_mut(13);
                    let (data_notag, tag_bytes) = data_and_tag.split_at_mut(
                        data_and_tag.len() + 1
                            - ::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);
                    }
                    final_len = data_notag.len();
                }
                data.drain(..Nonce::len());
                data.truncate(final_len);
                Ok(())
            }
        }
    }
    fn overhead(&self) -> usize {
        match self {
            Cipher::XChaCha20Poly1305(cipher) => {
                let cipher = CipherKind::XChaCha20Poly1305;
                cipher.nonce_len() + cipher.tag_len()
            }
        }
    }
    fn encrypt(
        &self,
        nonce: &Nonce,
        aad: AAD,
        data: &mut Data,
    ) -> Result<(), Error> {
        // No need to check for minimum buffer size since `Data` assures we
        // already went through that
        match self {
            Cipher::XChaCha20Poly1305(cipher) => {
                use ::chacha20poly1305::{
                    aead::generic_array::GenericArray, AeadInPlace,
                };
                // write nonce
                data.get_slice_full()[..Nonce::len()]
                    .copy_from_slice(nonce.as_bytes());

                // encrypt data
                match cipher.cipher.encrypt_in_place_detached(
                    nonce.as_bytes().into(),
                    aad.0,
                    data.get_slice(),
                ) {
                    Ok(tag) => {
                        // add tag
                        data.get_tag_slice().copy_from_slice(tag.as_slice());
                        Ok(())
                    }
                    Err(_) => Err(Error::Encrypt),
                };
            }
        }
        todo!()
    }
}

/// Receive only cipher
#[allow(missing_debug_implementations)]
pub struct CipherRecv(Cipher);

impl CipherRecv {
    /// Build a new Cipher
    pub fn new(kind: CipherKind, secret: Secret) -> Self {
        Self(Cipher::new(kind, secret))
    }
    /// Get the length of the nonce  for this cipher
    pub fn nonce_len(&self) -> usize {
        self.0.nonce_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: &mut VecDeque<u8>,
    ) -> Result<(), Error> {
        self.0.decrypt(aad, data)
    }
}

/// Allocate some data, with additional indexes to track
/// where nonce and tags are
#[derive(Debug, Clone)]
pub struct Data {
    data: Vec<u8>,
    skip_start: usize,
    skip_end: usize,
}

impl Data {
    /// Get the slice where you will write the actual data
    /// this will skip the actual nonce and AEAD tag and give you
    /// only the space for the data
    pub fn get_slice(&mut self) -> &mut [u8] {
        &mut self.data[self.skip_start..self.skip_end]
    }
    fn get_tag_slice(&mut self) -> &mut [u8] {
        let start = self.data.len() - self.skip_end;
        &mut self.data[start..]
    }
    fn get_slice_full(&mut self) -> &mut [u8] {
        &mut self.data
    }
    /// Consume the data and return the whole raw vector
    pub fn get_raw(self) -> Vec<u8> {
        self.data
    }
}

/// Send only cipher
#[allow(missing_debug_implementations)]
pub struct CipherSend {
    nonce: Nonce,
    cipher: Cipher,
}

impl CipherSend {
    /// Build a new Cipher
    pub fn new(kind: CipherKind, secret: Secret) -> Self {
        Self {
            nonce: Nonce::new(),
            cipher: Cipher::new(kind, secret),
        }
    }
    /// Allocate the memory for the data that will be encrypted
    pub fn make_data(&self, length: usize) -> Data {
        Data {
            data: Vec::with_capacity(length + self.cipher.overhead()),
            skip_start: self.cipher.nonce_len(),
            skip_end: self.cipher.tag_len(),
        }
    }
    /// Encrypt the given data
    pub fn encrypt(&mut self, aad: AAD, data: &mut Data) -> Result<(), Error> {
        self.cipher.encrypt(&self.nonce, aad, data)?;
        self.nonce.advance();
        Ok(())
    }
}

/// 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: For efficiency "Nonce" should become a reference.
//

#[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; 12],
}

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 {
    // FIXME: nonces should be random!
    /// Generate a new random Nonce
    pub fn new() -> Self {
        #[allow(unsafe_code)]
        unsafe {
            Self {
                // chosen by a fair dice roll
                // ahh, who am I kidding...
                num: NonceNum { high: 42, low: 69 },
            }
        }
    }
    /// Length of this nonce in bytes
    pub const fn len() -> usize {
        return 12;
    }
    /// 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; 12]) -> Self {
        #[allow(unsafe_code)]
        unsafe {
            Self { raw }
        }
    }
    /// Go to the next nonce
    pub fn advance(&mut self) {
        #[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;
            }
        }
    }
}