use aead::{AeadMutInPlace, KeyInit, KeySizeUser};
use chacha20poly1305::ChaCha20Poly1305;

use crate::crypt::symmetric::{
    data::{
        DecryptionBuffer, EncryptionBuffer, InboundEncryptedMessage, InboundPlainMessage, Iv,
        Nonce, OutboundEncryptedMessage, OutboundPlainMessage, PrefixedPayload,
    },
    make_aad, CiphersuiteAlgorithm, Decrypter, Encrypter, Error,
};

pub struct Chacha20Poly1305Algorithm;
pub struct Chacha20Poly1305Cipher(ChaCha20Poly1305, Iv);

pub const CHACHA20POLY1305_OVERHEAD: usize = 16;

pub static CHACHA20POLY1305: Chacha20Poly1305Algorithm = Chacha20Poly1305Algorithm;

impl CiphersuiteAlgorithm for Chacha20Poly1305Algorithm {
    fn encrypter(&self, key: &[u8], iv: &Iv) -> Result<Box<dyn Encrypter>, Error> {
        Ok(Box::new(Chacha20Poly1305Cipher(
            ChaCha20Poly1305::new_from_slice(key).map_err(|_| Error::InvalidKeySize)?,
            *iv,
        )))
    }

    fn decrypter(&self, key: &[u8], iv: &Iv) -> Result<Box<dyn Decrypter>, Error> {
        Ok(Box::new(Chacha20Poly1305Cipher(
            ChaCha20Poly1305::new_from_slice(key).map_err(|_| Error::InvalidKeySize)?,
            *iv,
        )))
    }

    fn key_size(&self) -> usize {
        ChaCha20Poly1305::key_size()
    }
}

impl Encrypter for Chacha20Poly1305Cipher {
    fn encrypt(
        &mut self,
        seq: u64,
        msg: OutboundPlainMessage<'_>,
    ) -> Result<OutboundEncryptedMessage, Error> {
        let total_len = msg.payload.len() + CHACHA20POLY1305_OVERHEAD;
        let mut payload = PrefixedPayload::with_capacity(total_len);

        payload.extend_from_slice(msg.payload);
        let nonce = chacha20poly1305::Nonce::from(Nonce::new(&self.1, seq).0);
        let aad = make_aad(total_len, seq);

        self.0
            .encrypt_in_place(&nonce, &aad, &mut EncryptionBuffer(&mut payload))
            .map_err(|_| Error::EncryptError)
            .map(|_| OutboundEncryptedMessage::new(payload))
    }
}

impl Decrypter for Chacha20Poly1305Cipher {
    fn decrypt<'a>(
        &mut self,
        seq: u64,
        mut msg: InboundEncryptedMessage<'a>,
    ) -> Result<InboundPlainMessage<'a>, Error> {
        let payload = &msg.payload;
        let nonce = chacha20poly1305::Nonce::from(Nonce::new(&self.1, seq).0);
        let aad = make_aad(payload.len(), seq);
        let payload = &mut msg.payload;

        self.0
            .decrypt_in_place(&nonce, &aad, &mut DecryptionBuffer(payload))
            .map_err(|_| Error::DecryptError)?;

        Ok(msg.into_plain_message())
    }
}

#[cfg(test)]
mod tests {
    use crate::crypt::symmetric::{
        chacha20poly1305::CHACHA20POLY1305,
        data::{InboundEncryptedMessage, Iv, OutboundPlainMessage},
        CiphersuiteAlgorithm,
    };

    #[test]
    fn test_encrypt_decrypt() {
        let plaintext = b"This is a message";
        let key = b"11112222333344441111222233334444";
        let iv = Iv::from(*b"1111111111111111");

        let mut encrypter = CHACHA20POLY1305.encrypter(key, &iv).unwrap();
        let mut decrypter = CHACHA20POLY1305.decrypter(key, &iv).unwrap();

        let out_plain = OutboundPlainMessage::new(plaintext);
        let mut out_encrypted = encrypter.encrypt(1234, out_plain).unwrap();
        let in_encrypted = InboundEncryptedMessage::from(&mut out_encrypted.payload[..]);
        let in_plain = decrypter.decrypt(1234, in_encrypted).unwrap();

        assert_eq!(in_plain.payload, plaintext);
    }
}
