use std::ops::{Deref, DerefMut};

use crypto_common::typenum;

use crate::crypt::V1_CLIENT_RANDOM_LEN;

pub type NonceSize = typenum::U12;
pub const NONCE_SIZE: usize = 12;

pub struct Nonce(pub [u8; NONCE_SIZE]);
#[derive(Clone, Copy)]
pub struct Iv(pub [u8; V1_CLIENT_RANDOM_LEN]);

pub struct BorrowedPayload<'a>(&'a mut [u8]);
pub struct PrefixedPayload(Vec<u8>);

pub struct EncryptionBuffer<'a>(pub(crate) &'a mut PrefixedPayload);
pub struct DecryptionBuffer<'a, 'p>(pub(crate) &'a mut BorrowedPayload<'p>);

pub struct InboundPlainMessage<'a> {
    pub payload: &'a [u8],
}

pub struct InboundEncryptedMessage<'a> {
    pub payload: BorrowedPayload<'a>,
}

pub struct OutboundPlainMessage<'a> {
    pub payload: &'a [u8],
}

pub struct OutboundEncryptedMessage {
    pub payload: PrefixedPayload,
}

impl From<[u8; V1_CLIENT_RANDOM_LEN]> for Iv {
    fn from(value: [u8; V1_CLIENT_RANDOM_LEN]) -> Self {
        Self(value)
    }
}

impl From<[u8; NONCE_SIZE]> for Nonce {
    fn from(value: [u8; NONCE_SIZE]) -> Self {
        Self(value)
    }
}

impl Nonce {
    pub fn new(iv: &Iv, seq: u64) -> Self {
        let mut data = [0; NONCE_SIZE];
        data[0..8].copy_from_slice(&seq.to_be_bytes());
        data.iter_mut().zip(iv.0.iter()).for_each(|(data, iv)| {
            *data ^= *iv;
        });
        Self(data)
    }
}

impl<'a> InboundEncryptedMessage<'a> {
    pub fn into_plain_message(self) -> InboundPlainMessage<'a> {
        InboundPlainMessage {
            payload: self.payload.into_inner(),
        }
    }
}

impl<'a> From<&'a mut [u8]> for InboundEncryptedMessage<'a> {
    fn from(value: &'a mut [u8]) -> Self {
        Self {
            payload: BorrowedPayload(value),
        }
    }
}

impl<'a> OutboundPlainMessage<'a> {
    pub fn new(payload: &'a [u8]) -> Self {
        Self { payload }
    }
}

impl OutboundEncryptedMessage {
    pub fn new(payload: PrefixedPayload) -> Self {
        Self { payload }
    }
}

impl<'a> BorrowedPayload<'a> {
    pub fn new(payload: &'a mut [u8]) -> Self {
        Self(payload)
    }

    pub fn into_inner(self) -> &'a mut [u8] {
        self.0
    }

    pub fn truncate(&mut self, len: usize) {
        if len >= self.len() {
            return;
        }

        self.0 = core::mem::take(&mut self.0).split_at_mut(len).0;
    }
}

impl Deref for BorrowedPayload<'_> {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        self.0
    }
}

impl DerefMut for BorrowedPayload<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.0
    }
}

impl PrefixedPayload {
    pub fn with_capacity(capacity: usize) -> Self {
        Self(Vec::with_capacity(capacity))
    }

    pub fn extend_from_slice(&mut self, slice: &[u8]) {
        self.0.extend_from_slice(slice);
    }

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

impl Deref for PrefixedPayload {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for PrefixedPayload {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

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

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

impl aead::Buffer for EncryptionBuffer<'_> {
    fn len(&self) -> usize {
        self.0.len()
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

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

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

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

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

impl aead::Buffer for DecryptionBuffer<'_, '_> {
    fn len(&self) -> usize {
        self.0.len()
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

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

    fn extend_from_slice(&mut self, _other: &[u8]) -> aead::Result<()> {
        todo!("Not used in decryption")
    }
}
