use std::marker::PhantomData;

use aead::{KeyInit, KeySizeUser};
use aes::cipher::{BlockCipher, BlockEncrypt};
use aes_gcm::{aead::AeadMutInPlace, AesGcm};
use crypto_common::{typenum, BlockSizeUser};

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

pub struct AesGcmAlgorithm<Aes>(PhantomData<Aes>);
pub struct AesGcmCipher<Aes>(AesGcm<Aes, NonceSize>, Iv);

pub const AES_GCM_OVERHEAD: usize = 16;

pub static AES128GCM: AesGcmAlgorithm<aes::Aes128> = AesGcmAlgorithm(PhantomData);
pub static AES256GCM: AesGcmAlgorithm<aes::Aes256> = AesGcmAlgorithm(PhantomData);

impl<Aes> CiphersuiteAlgorithm for AesGcmAlgorithm<Aes>
where
    Aes: BlockCipher + BlockEncrypt + BlockSizeUser<BlockSize = typenum::U16> + KeyInit + 'static,
    AesGcm<Aes, NonceSize>: KeySizeUser + KeyInit,
{
    fn encrypter(&self, key: &[u8], iv: &Iv) -> Result<Box<dyn Encrypter>, Error> {
        Ok(Box::new(AesGcmCipher::<Aes>(
            AesGcm::new_from_slice(key).map_err(|_| Error::InvalidKeySize)?,
            *iv,
        )))
    }

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

    fn key_size(&self) -> usize {
        AesGcm::<Aes, NonceSize>::key_size()
    }
}

impl<Aes> Encrypter for AesGcmCipher<Aes>
where
    Aes: BlockCipher + BlockEncrypt + BlockSizeUser<BlockSize = typenum::U16>,
{
    fn encrypt(
        &mut self,
        seq: u64,
        msg: OutboundPlainMessage<'_>,
    ) -> Result<OutboundEncryptedMessage, Error> {
        let total_len = msg.payload.len() + AES_GCM_OVERHEAD;
        let mut payload = PrefixedPayload::with_capacity(total_len);

        payload.extend_from_slice(msg.payload);
        let nonce = aes_gcm::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<Aes> Decrypter for AesGcmCipher<Aes>
where
    Aes: BlockCipher + BlockEncrypt + BlockSizeUser<BlockSize = typenum::U16>,
{
    fn decrypt<'a>(
        &mut self,
        seq: u64,
        mut msg: InboundEncryptedMessage<'a>,
    ) -> Result<InboundPlainMessage<'a>, Error> {
        let payload = &msg.payload;
        let nonce = aes_gcm::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::{
        aes_gcm::{AES128GCM, AES256GCM},
        data::{InboundEncryptedMessage, Iv, OutboundPlainMessage},
        CiphersuiteAlgorithm,
    };

    #[test]
    fn test_aes_gcm_encrypt_decrypt() {
        const ALGORITHMS: &[(&dyn CiphersuiteAlgorithm, &[u8])] = &[
            (&AES128GCM, b"1234123412341234"),
            (&AES256GCM, b"12341234123412345678567856785678"),
        ];
        let plaintext = b"This is a message";
        let iv = Iv::from(*b"1111111111111111");

        for (aes, key) in ALGORITHMS {
            let mut encrypter = aes.encrypter(key, &iv).unwrap();
            let mut decrypter = aes.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);
        }
    }
}
