//! Implementation of the [`AeadChaCha20Poly1305`] trait based on the [`libcrux_chacha20poly1305`] crate.

use rosenpass_cipher_traits::algorithms::AeadChaCha20Poly1305;
use rosenpass_cipher_traits::primitives::{Aead, AeadError};

pub use rosenpass_cipher_traits::algorithms::aead_chacha20poly1305::{KEY_LEN, NONCE_LEN, TAG_LEN};

/// An implementation of the ChaCha20Poly1305 AEAD based on libcrux
pub struct ChaCha20Poly1305;

impl Aead<KEY_LEN, NONCE_LEN, TAG_LEN> for ChaCha20Poly1305 {
    fn encrypt(
        &self,
        ciphertext: &mut [u8],
        key: &[u8; KEY_LEN],
        nonce: &[u8; NONCE_LEN],
        ad: &[u8],
        plaintext: &[u8],
    ) -> Result<(), AeadError> {
        let (ctxt, tag) = libcrux_chacha20poly1305::encrypt(key, plaintext, ciphertext, ad, nonce)
            .map_err(|_| AeadError::InternalError)?;

        // return an error of the destination buffer is longer than expected
        // because the caller wouldn't know where the end is
        if ctxt.len() + tag.len() != ciphertext.len() {
            return Err(AeadError::InternalError);
        }

        Ok(())
    }

    fn decrypt(
        &self,
        plaintext: &mut [u8],
        key: &[u8; KEY_LEN],
        nonce: &[u8; NONCE_LEN],
        ad: &[u8],
        ciphertext: &[u8],
    ) -> Result<(), AeadError> {
        let ptxt = libcrux_chacha20poly1305::decrypt(key, plaintext, ciphertext, ad, nonce)
            .map_err(|_| AeadError::DecryptError)?;

        // return an error of the destination buffer is longer than expected
        // because the caller wouldn't know where the end is
        if ptxt.len() != plaintext.len() {
            return Err(AeadError::DecryptError);
        }

        Ok(())
    }
}

impl AeadChaCha20Poly1305 for ChaCha20Poly1305 {}

/// The idea of these tests is to check that the above implemenatation behaves, by and large, the
/// same as the one from the old libcrux and the one from RustCrypto. You can consider them janky,
/// self-rolled property-based tests.
#[cfg(test)]
mod equivalence_tests {
    use super::*;
    use rand::RngCore;

    #[test]
    fn proptest_equivalence_libcrux_rustcrypto() {
        use crate::subtle::rust_crypto::chacha20poly1305_ietf::ChaCha20Poly1305 as RustCryptoChaCha20Poly1305;
        let ptxts: [&[u8]; 3] = [
            b"".as_slice(),
            b"test".as_slice(),
            b"abcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcd",
        ];
        let mut key = [0; KEY_LEN];
        let mut rng = rand::thread_rng();

        let mut ctxt_left = [0; 64 + TAG_LEN];
        let mut ctxt_right = [0; 64 + TAG_LEN];

        let mut ptxt_left = [0; 64];
        let mut ptxt_right = [0; 64];

        let nonce = [0; NONCE_LEN];
        let ad = b"";

        for ptxt in ptxts {
            for _ in 0..1000 {
                rng.fill_bytes(&mut key);
                let ctxt_left = &mut ctxt_left[..ptxt.len() + TAG_LEN];
                let ctxt_right = &mut ctxt_right[..ptxt.len() + TAG_LEN];

                let ptxt_left = &mut ptxt_left[..ptxt.len()];
                let ptxt_right = &mut ptxt_right[..ptxt.len()];

                RustCryptoChaCha20Poly1305
                    .encrypt(ctxt_left, &key, &nonce, ad, ptxt)
                    .unwrap();
                ChaCha20Poly1305
                    .encrypt(ctxt_right, &key, &nonce, ad, ptxt)
                    .unwrap();

                assert_eq!(ctxt_left, ctxt_right);

                RustCryptoChaCha20Poly1305
                    .decrypt(ptxt_left, &key, &nonce, ad, ctxt_left)
                    .unwrap();
                ChaCha20Poly1305
                    .decrypt(ptxt_right, &key, &nonce, ad, ctxt_right)
                    .unwrap();

                assert_eq!(ptxt_left, ptxt);
                assert_eq!(ptxt_right, ptxt);
            }
        }
    }

    #[test]
    #[cfg(feature = "experiment_libcrux_chachapoly_test")]
    fn proptest_equivalence_libcrux_old_new() {
        let ptxts: [&[u8]; 3] = [
            b"".as_slice(),
            b"test".as_slice(),
            b"abcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcdabcd",
        ];
        let mut key = [0; KEY_LEN];
        let mut rng = rand::thread_rng();

        let mut ctxt_left = [0; 64 + TAG_LEN];
        let mut ctxt_right = [0; 64 + TAG_LEN];

        let mut ptxt_left = [0; 64];
        let mut ptxt_right = [0; 64];

        let nonce = [0; NONCE_LEN];
        let ad = b"";

        for ptxt in ptxts {
            for _ in 0..1000 {
                rng.fill_bytes(&mut key);
                let ctxt_left = &mut ctxt_left[..ptxt.len() + TAG_LEN];
                let ctxt_right = &mut ctxt_right[..ptxt.len() + TAG_LEN];

                let ptxt_left = &mut ptxt_left[..ptxt.len()];
                let ptxt_right = &mut ptxt_right[..ptxt.len()];

                encrypt(ctxt_left, &key, &nonce, ad, ptxt).unwrap();
                ChaCha20Poly1305
                    .encrypt(ctxt_right, &key, &nonce, ad, ptxt)
                    .unwrap();

                assert_eq!(ctxt_left, ctxt_right);

                decrypt(ptxt_left, &key, &nonce, ad, ctxt_left).unwrap();
                ChaCha20Poly1305
                    .decrypt(ptxt_right, &key, &nonce, ad, ctxt_right)
                    .unwrap();

                assert_eq!(ptxt_left, ptxt);
                assert_eq!(ptxt_right, ptxt);
            }
        }

        // The old libcrux functions:

        // The functions below are from the old libcrux backend. I am keeping them around so we can
        // check if they behave the same.
        use rosenpass_to::ops::copy_slice;
        use rosenpass_to::To;
        use zeroize::Zeroize;

        /// Encrypts using ChaCha20Poly1305 as implemented in [libcrux](https://github.com/cryspen/libcrux).
        /// Key and nonce MUST be chosen (pseudo-)randomly. The `key` slice MUST have a length of
        /// [KEY_LEN]. The `nonce` slice MUST have a length of [NONCE_LEN]. The last [TAG_LEN] bytes
        /// written in `ciphertext` are the tag guaranteeing integrity. `ciphertext` MUST have a capacity of
        /// `plaintext.len()` + [TAG_LEN].
        ///  
        /// # Examples
        ///```rust
        /// # use rosenpass_ciphers::subtle::chacha20poly1305_ietf_libcrux::{encrypt, TAG_LEN, KEY_LEN, NONCE_LEN};
        ///
        /// const PLAINTEXT_LEN: usize = 43;
        /// let plaintext = "post-quantum cryptography is very important".as_bytes();
        /// assert_eq!(PLAINTEXT_LEN, plaintext.len());
        /// let key: &[u8] = &[0u8; KEY_LEN]; // THIS IS NOT A SECURE KEY
        /// let nonce: &[u8] = &[0u8; NONCE_LEN]; // THIS IS NOT A SECURE NONCE
        /// let additional_data: &[u8] = "the encrypted message is very important".as_bytes();
        /// let mut ciphertext_buffer = [0u8; PLAINTEXT_LEN + TAG_LEN];
        ///
        /// let res: anyhow::Result<()> = encrypt(&mut ciphertext_buffer, key, nonce, additional_data, plaintext);
        /// assert!(res.is_ok());
        /// # let expected_ciphertext: &[u8] = &[239, 104, 148, 202, 120, 32, 77, 27, 246, 206, 226, 17,
        /// # 83, 78, 122, 116, 187, 123, 70, 199, 58, 130, 21, 1, 107, 230, 58, 77, 18, 152, 31, 159, 80,
        /// # 151, 72, 27, 236, 137, 60, 55, 180, 31, 71, 97, 199, 12, 60, 155, 70, 221, 225, 110, 132, 191,
        /// # 8, 114, 85, 4, 25];
        /// # assert_eq!(expected_ciphertext, &ciphertext_buffer);
        ///```
        ///
        #[inline]
        pub fn encrypt(
            ciphertext: &mut [u8],
            key: &[u8],
            nonce: &[u8],
            ad: &[u8],
            plaintext: &[u8],
        ) -> anyhow::Result<()> {
            let (ciphertext, mac) = ciphertext.split_at_mut(ciphertext.len() - TAG_LEN);

            use libcrux::aead as C;
            let crux_key = C::Key::Chacha20Poly1305(C::Chacha20Key(key.try_into().unwrap()));
            let crux_iv = C::Iv(nonce.try_into().unwrap());

            copy_slice(plaintext).to(ciphertext);
            let crux_tag = libcrux::aead::encrypt(&crux_key, ciphertext, crux_iv, ad).unwrap();
            copy_slice(crux_tag.as_ref()).to(mac);

            match crux_key {
                C::Key::Chacha20Poly1305(mut k) => k.0.zeroize(),
                _ => panic!(),
            }

            Ok(())
        }

        /// Decrypts a `ciphertext` and verifies the integrity of the `ciphertext` and the additional data
        /// `ad`. using ChaCha20Poly1305 as implemented in [libcrux](https://github.com/cryspen/libcrux).
        ///
        /// The `key` slice MUST have a length of [KEY_LEN]. The `nonce` slice MUST have a length of
        /// [NONCE_LEN]. The plaintext buffer must have a capacity of `ciphertext.len()` - [TAG_LEN].
        ///
        /// # Examples
        ///```rust
        /// # use rosenpass_ciphers::subtle::chacha20poly1305_ietf_libcrux::{decrypt, TAG_LEN, KEY_LEN, NONCE_LEN};
        /// let ciphertext: &[u8] = &[239, 104, 148, 202, 120, 32, 77, 27, 246, 206, 226, 17,
        /// 83, 78, 122, 116, 187, 123, 70, 199, 58, 130, 21, 1, 107, 230, 58, 77, 18, 152, 31, 159, 80,
        /// 151, 72, 27, 236, 137, 60, 55, 180, 31, 71, 97, 199, 12, 60, 155, 70, 221, 225, 110, 132, 191,
        /// 8, 114, 85, 4, 25]; // this is the ciphertext generated by the example for the encryption
        /// const PLAINTEXT_LEN: usize = 43;
        /// assert_eq!(PLAINTEXT_LEN + TAG_LEN, ciphertext.len());
        ///
        /// let key: &[u8] = &[0u8; KEY_LEN]; // THIS IS NOT A SECURE KEY
        /// let nonce: &[u8] = &[0u8; NONCE_LEN]; // THIS IS NOT A SECURE NONCE
        /// let additional_data: &[u8] = "the encrypted message is very important".as_bytes();
        /// let mut plaintext_buffer = [0u8; PLAINTEXT_LEN];
        ///
        /// let res: anyhow::Result<()> = decrypt(&mut plaintext_buffer, key, nonce, additional_data, ciphertext);
        /// assert!(res.is_ok());
        /// let expected_plaintext = "post-quantum cryptography is very important".as_bytes();
        /// assert_eq!(expected_plaintext, plaintext_buffer);
        ///
        ///```
        #[inline]
        pub fn decrypt(
            plaintext: &mut [u8],
            key: &[u8],
            nonce: &[u8],
            ad: &[u8],
            ciphertext: &[u8],
        ) -> anyhow::Result<()> {
            let (ciphertext, mac) = ciphertext.split_at(ciphertext.len() - TAG_LEN);

            use libcrux::aead as C;
            let crux_key = C::Key::Chacha20Poly1305(C::Chacha20Key(key.try_into().unwrap()));
            let crux_iv = C::Iv(nonce.try_into().unwrap());
            let crux_tag = C::Tag::from_slice(mac).unwrap();

            copy_slice(ciphertext).to(plaintext);
            libcrux::aead::decrypt(&crux_key, plaintext, crux_iv, ad, &crux_tag).unwrap();

            match crux_key {
                C::Key::Chacha20Poly1305(mut k) => k.0.zeroize(),
                _ => panic!(),
            }

            Ok(())
        }
    }
}
