use std::marker::PhantomData;

use aead::{generic_array::GenericArray, AeadInPlace, KeyInit};
use rustls::{
    crypto::cipher::{
        make_tls12_aad, make_tls13_aad, AeadKey, BorrowedPayload, InboundOpaqueMessage,
        InboundPlainMessage, Iv, KeyBlockShape, MessageDecrypter, MessageEncrypter, Nonce,
        OutboundOpaqueMessage, OutboundPlainMessage, PrefixedPayload, Tls12AeadAlgorithm,
        Tls13AeadAlgorithm, UnsupportedOperationError,
    },
    ConnectionTrafficSecrets, ContentType, ProtocolVersion,
};

use crate::traits::{AeadDataInput, TlsExtractKeys};

mod imp;

pub use imp::*;

pub(crate) struct Tls13Cipher<Aead>(pub Aead, pub Iv);
pub(crate) struct Tls12Cipher<Aead>(pub Aead, pub Iv);

pub struct TlsAeadCipherAlgorithm<Aead>(PhantomData<Aead>);

struct EncryptBufferAdapter<'a>(&'a mut PrefixedPayload);
struct DecryptBufferAdapter<'a, 'p>(&'a mut BorrowedPayload<'p>);

impl AsRef<[u8]> for EncryptBufferAdapter<'_> {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl AsMut<[u8]> for EncryptBufferAdapter<'_> {
    fn as_mut(&mut self) -> &mut [u8] {
        self.0.as_mut()
    }
}

impl aead::Buffer for EncryptBufferAdapter<'_> {
    fn extend_from_slice(&mut self, other: &[u8]) -> aead::Result<()> {
        self.0.extend_from_slice(other);
        Ok(())
    }

    fn truncate(&mut self, len: usize) {
        self.0.truncate(len);
    }
}

impl AsRef<[u8]> for DecryptBufferAdapter<'_, '_> {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl AsMut<[u8]> for DecryptBufferAdapter<'_, '_> {
    fn as_mut(&mut self) -> &mut [u8] {
        self.0.as_mut()
    }
}

impl aead::Buffer for DecryptBufferAdapter<'_, '_> {
    fn extend_from_slice(&mut self, other: &[u8]) -> aead::Result<()> {
        let _ = other;
        unreachable!("not used by `AeadInPlace::decrypt_in_place`")
    }

    fn truncate(&mut self, len: usize) {
        self.0.truncate(len);
    }
}

impl<Aead> MessageEncrypter for Tls13Cipher<Aead>
where
    Aead: Send + Sync + AeadInPlace + AeadDataInput + 'static,
    GenericArray<u8, Aead::NonceSize>: From<[u8; 12]>,
{
    fn encrypt(
        &mut self,
        msg: OutboundPlainMessage<'_>,
        seq: u64,
    ) -> Result<OutboundOpaqueMessage, rustls::Error> {
        let total_len = self.encrypted_payload_len(msg.payload.len());
        let mut payload = PrefixedPayload::with_capacity(total_len);

        payload.extend_from_chunks(&msg.payload);
        payload.extend_from_slice(&msg.typ.to_array());
        let nonce = aead::Nonce::<Aead>::from(Nonce::new(&self.1, seq).0);
        let aad = make_tls13_aad(total_len);

        self.0
            .encrypt_in_place(&nonce, &aad, &mut EncryptBufferAdapter(&mut payload))
            .map_err(|_| rustls::Error::EncryptError)
            .map(|_| {
                OutboundOpaqueMessage::new(
                    ContentType::ApplicationData,
                    ProtocolVersion::TLSv1_2,
                    payload,
                )
            })
    }

    fn encrypted_payload_len(&self, payload_len: usize) -> usize {
        payload_len + Aead::ENCRYPTED_OVERHEAD + 1
    }
}

impl<Aead> MessageDecrypter for Tls13Cipher<Aead>
where
    Aead: Send + Sync + AeadInPlace + AeadDataInput + 'static,
    GenericArray<u8, Aead::NonceSize>: From<[u8; 12]>,
{
    fn decrypt<'a>(
        &mut self,
        mut msg: InboundOpaqueMessage<'a>,
        seq: u64,
    ) -> Result<InboundPlainMessage<'a>, rustls::Error> {
        let payload = &mut msg.payload;
        let nonce = aead::Nonce::<Aead>::from(Nonce::new(&self.1, seq).0);
        let aad = make_tls13_aad(payload.len());

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

        msg.into_tls13_unpadded_message()
    }
}

impl<Aead> MessageEncrypter for Tls12Cipher<Aead>
where
    Aead: Send + Sync + AeadInPlace + AeadDataInput + 'static,
    GenericArray<u8, Aead::NonceSize>: From<[u8; 12]>,
{
    fn encrypt(
        &mut self,
        msg: OutboundPlainMessage<'_>,
        seq: u64,
    ) -> Result<OutboundOpaqueMessage, rustls::Error> {
        let total_len = self.encrypted_payload_len(msg.payload.len());
        let mut payload = PrefixedPayload::with_capacity(total_len);

        payload.extend_from_chunks(&msg.payload);
        let nonce = aead::Nonce::<Aead>::from(Nonce::new(&self.1, 0).0);
        let aad = make_tls12_aad(seq, msg.typ, msg.version, msg.payload.len());

        self.0
            .encrypt_in_place(&nonce, &aad, &mut EncryptBufferAdapter(&mut payload))
            .map_err(|_| rustls::Error::EncryptError)
            .map(|_| OutboundOpaqueMessage::new(msg.typ, msg.version, payload))
    }

    fn encrypted_payload_len(&self, payload_len: usize) -> usize {
        payload_len + Aead::ENCRYPTED_OVERHEAD
    }
}

impl<Aead> MessageDecrypter for Tls12Cipher<Aead>
where
    Aead: Send + Sync + AeadInPlace + AeadDataInput + 'static,
    GenericArray<u8, Aead::NonceSize>: From<[u8; 12]>,
{
    fn decrypt<'a>(
        &mut self,
        mut msg: InboundOpaqueMessage<'a>,
        seq: u64,
    ) -> Result<InboundPlainMessage<'a>, rustls::Error> {
        let payload = &msg.payload;
        let nonce = aead::Nonce::<Aead>::from(Nonce::new(&self.1, seq).0);
        let aad = make_tls12_aad(
            seq,
            msg.typ,
            msg.version,
            payload.len() - Aead::ENCRYPTED_OVERHEAD,
        );

        let payload = &mut msg.payload;
        self.0
            .decrypt_in_place(&nonce, &aad, &mut DecryptBufferAdapter(payload))
            .map_err(|_| rustls::Error::DecryptError)?;

        Ok(msg.into_plain_message())
    }
}

impl<Aead> Tls13AeadAlgorithm for TlsAeadCipherAlgorithm<Aead>
where
    Aead: Send + Sync + TlsExtractKeys + KeyInit + 'static,
    Tls13Cipher<Aead>: MessageEncrypter + MessageDecrypter,
{
    fn encrypter(&self, key: AeadKey, iv: Iv) -> Box<dyn MessageEncrypter> {
        let aead = Aead::new_from_slice(key.as_ref()).unwrap();
        let cipher = Tls13Cipher::<Aead>(aead, iv);
        Box::new(cipher)
    }

    fn decrypter(&self, key: AeadKey, iv: Iv) -> Box<dyn MessageDecrypter> {
        let aead = Aead::new_from_slice(key.as_ref()).unwrap();
        let cipher = Tls13Cipher::<Aead>(aead, iv);
        Box::new(cipher)
    }

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

    fn extract_keys(
        &self,
        key: AeadKey,
        iv: Iv,
    ) -> Result<ConnectionTrafficSecrets, UnsupportedOperationError> {
        Ok(Aead::tls13_wrap_keys(key, iv))
    }
}

impl<Aead> Tls12AeadAlgorithm for TlsAeadCipherAlgorithm<Aead>
where
    Aead: Send + Sync + TlsExtractKeys + KeyInit + 'static,
    Tls12Cipher<Aead>: MessageEncrypter + MessageDecrypter,
{
    fn encrypter(&self, key: AeadKey, iv: &[u8], _extra: &[u8]) -> Box<dyn MessageEncrypter> {
        let aead = Aead::new_from_slice(key.as_ref()).unwrap();
        let cipher = Tls12Cipher::<Aead>(aead, Iv::copy(iv));
        Box::new(cipher)
    }

    fn decrypter(&self, key: AeadKey, iv: &[u8]) -> Box<dyn MessageDecrypter> {
        let aead = Aead::new_from_slice(key.as_ref()).unwrap();
        let cipher = Tls12Cipher::<Aead>(aead, Iv::copy(iv));
        Box::new(cipher)
    }

    fn key_block_shape(&self) -> KeyBlockShape {
        KeyBlockShape {
            enc_key_len: Aead::key_size(),
            fixed_iv_len: 12,
            explicit_nonce_len: 0,
        }
    }

    fn extract_keys(
        &self,
        key: AeadKey,
        iv: &[u8],
        explicit: &[u8],
    ) -> Result<ConnectionTrafficSecrets, UnsupportedOperationError> {
        Ok(Aead::tls12_wrap_keys(key, iv, explicit))
    }
}
