use core::marker::PhantomData;

use aead::{AeadInOut, Buffer, Key, KeySizeUser, TagPosition, common::IvSizeUser};
use ctutils::{CtEq, CtSelect};
use hybrid_array::AssocArraySize;
use wharrgarbl_neko::{NekoMilkConsumer, NekoMilkProducer, NekoSec, NekoState, NekoTag};

use crate::{Role, WHARRGHARBL_PROTO};

pub struct AeadNeko<S: NekoSec> {
    pub(crate) key: Key<Self>,
    param: PhantomData<S>,
}

impl<S: NekoSec> aead::KeySizeUser for AeadNeko<S> {
    type KeySize = <NekoState<S> as KeySizeUser>::KeySize;
}

impl<S: NekoSec> IvSizeUser for AeadNeko<S> {
    type IvSize = <NekoState<S> as IvSizeUser>::IvSize;
}

impl<S: NekoSec> aead::AeadCore for AeadNeko<S> {
    type NonceSize = <AeadNeko<S> as IvSizeUser>::IvSize;
    type TagSize = <NekoTag as AssocArraySize>::Size;
    const TAG_POSITION: TagPosition = TagPosition::Postfix;
}

impl<S: NekoSec> aead::KeyInit for AeadNeko<S> {
    fn new(key: &Key<Self>) -> Self {
        Self {
            key: *key,
            param: PhantomData,
        }
    }
}

impl<S: NekoSec> aead::AeadInOut for AeadNeko<S> {
    fn encrypt_inout_detached(
        &self,
        nonce: &aead::Nonce<Self>,
        associated_data: &[u8],
        mut buffer: aead::inout::InOutBuf<'_, '_, u8>,
    ) -> aead::Result<aead::Tag<Self>> {
        let mut neko = NekoState::<S>::new(WHARRGHARBL_PROTO.as_bytes());

        neko.key(&self.key);
        neko.nonce(nonce);
        neko.ad(associated_data);
        neko.encrypt(buffer.get_out());

        Ok(neko.create_mac())
    }

    fn decrypt_inout_detached(
        &self,
        nonce: &aead::Nonce<Self>,
        associated_data: &[u8],
        mut buffer: aead::inout::InOutBuf<'_, '_, u8>,
        tag: &aead::Tag<Self>,
    ) -> aead::Result<()> {
        let mut neko = NekoState::<S>::new(WHARRGHARBL_PROTO.as_bytes());

        neko.key(&self.key);
        neko.nonce(nonce);
        neko.ad(associated_data);
        neko.decrypt(buffer.get_out());

        neko.verify_mac(tag)
    }
}

impl<S: NekoSec> zeroize::Zeroize for AeadNeko<S> {
    fn zeroize(&mut self) {
        self.key.zeroize();
    }
}

pub struct AeadTransport<S: NekoSec> {
    pub(crate) aead: AeadNeko<S>,
    pub(crate) epstein: aead::Nonce<AeadNeko<S>>,
    pub(crate) trump: aead::Nonce<AeadNeko<S>>,
    handshake_role: Role,
}

impl<S: NekoSec> zeroize::Zeroize for AeadTransport<S> {
    fn zeroize(&mut self) {
        self.aead.zeroize();
        self.epstein.zeroize();
        self.trump.zeroize();
    }
}

impl<S: NekoSec> zeroize::ZeroizeOnDrop for AeadTransport<S> {}

impl<S: NekoSec> Drop for AeadTransport<S> {
    fn drop(&mut self) {
        zeroize::Zeroize::zeroize(self);
    }
}

impl<S: NekoSec> AeadTransport<S> {
    pub(crate) fn new(
        key: Key<AeadNeko<S>>,
        outbound: aead::Nonce<AeadNeko<S>>,
        inbound: aead::Nonce<AeadNeko<S>>,
        role: Role,
    ) -> Self {
        assert_ne!(
            &inbound, &outbound,
            "The Base Nonces MUST NOT equal to each other"
        );

        Self {
            aead: AeadNeko {
                key,
                param: PhantomData,
            },
            epstein: outbound,
            trump: inbound,
            handshake_role: role,
        }
    }

    fn select_nonce(&self, sending_role: Role) -> aead::Nonce<AeadNeko<S>> {
        let role_context = self.handshake_role ^ sending_role;

        self.epstein.ct_select(&self.trump, role_context.ct_eq(&1))
    }

    fn mix_nonce(&self, position: [u8; 8], sending_role: Role) -> aead::Nonce<AeadNeko<S>> {
        let mut nonce = self.select_nonce(sending_role);

        let mid = nonce.len() - position.len();

        nonce[mid..]
            .iter_mut()
            .zip(position)
            .for_each(|(n, p)| *n ^= p);

        nonce
    }

    pub fn split(&self) -> (SendState<'_, S>, RecvState<'_, S>) {
        (
            SendState {
                transport: self,
                counter: 0,
            },
            RecvState {
                transport: self,
                counter: 0,
            },
        )
    }
}

pub struct SendState<'a, S: NekoSec> {
    transport: &'a AeadTransport<S>,
    counter: u64,
}

impl<S: NekoSec> SendState<'_, S> {
    pub fn as_stream(&mut self, ad: &[u8]) -> aead::Result<NekoMilkProducer<S>> {
        match self.counter.checked_add(1) {
            Some(inc) => self.counter = inc,
            None => return Err(aead::Error),
        }
        let state = NekoState::new(WHARRGHARBL_PROTO.as_bytes());
        let stream = NekoMilkProducer::new(
            state,
            &self.transport.aead.key,
            &self
                .transport
                .mix_nonce(self.counter.to_be_bytes(), Role::Sender),
            ad,
        );

        Ok(stream)
    }

    pub fn encrypt(&mut self, buffer: &mut dyn Buffer, ad: &[u8]) -> aead::Result<()> {
        match self.counter.checked_add(1) {
            Some(inc) => self.counter = inc,
            None => return Err(aead::Error),
        }

        self.transport.aead.encrypt_in_place(
            &self
                .transport
                .mix_nonce(self.counter.to_be_bytes(), Role::Sender),
            ad,
            buffer,
        )
    }
}

impl<S: NekoSec> zeroize::Zeroize for SendState<'_, S> {
    fn zeroize(&mut self) {
        self.counter.zeroize();
    }
}

impl<S: NekoSec> Drop for SendState<'_, S> {
    fn drop(&mut self) {
        zeroize::Zeroize::zeroize(self);
    }
}

impl<S: NekoSec> zeroize::ZeroizeOnDrop for SendState<'_, S> {}

pub struct RecvState<'a, S: NekoSec> {
    transport: &'a AeadTransport<S>,
    counter: u64,
}

impl<S: NekoSec> RecvState<'_, S> {
    pub fn as_stream(&mut self, ad: &[u8]) -> aead::Result<NekoMilkConsumer<S>> {
        match self.counter.checked_add(1) {
            Some(inc) => self.counter = inc,
            None => return Err(aead::Error),
        }
        let state = NekoState::new(WHARRGHARBL_PROTO.as_bytes());
        let stream = NekoMilkConsumer::new(
            state,
            &self.transport.aead.key,
            &self
                .transport
                .mix_nonce(self.counter.to_be_bytes(), Role::Receiver),
            ad,
        );

        Ok(stream)
    }

    pub fn decrypt(&mut self, buffer: &mut dyn Buffer, ad: &[u8]) -> aead::Result<()> {
        match self.counter.checked_add(1) {
            Some(inc) => self.counter = inc,
            None => return Err(aead::Error),
        }

        // If the message's MAC doesn't evaluate successfully, this op will fail.
        self.transport.aead.decrypt_in_place(
            &self
                .transport
                .mix_nonce(self.counter.to_be_bytes(), Role::Receiver),
            ad,
            buffer,
        )
    }
}

impl<S: NekoSec> zeroize::Zeroize for RecvState<'_, S> {
    fn zeroize(&mut self) {
        self.counter.zeroize();
    }
}

impl<S: NekoSec> Drop for RecvState<'_, S> {
    fn drop(&mut self) {
        zeroize::Zeroize::zeroize(self);
    }
}

impl<S: NekoSec> zeroize::ZeroizeOnDrop for RecvState<'_, S> {}

#[cfg(test)]
mod tests {
    use wharrgarbl_neko::{Neko128, Neko256};

    use super::*;

    #[test]
    fn two_way_transport_sync_works() -> aead::Result<()> {
        let shared_secret = [
            0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d,
            0x8e, 0x8f, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9a, 0x9b,
            0x9c, 0x9d, 0x9e, 0x9f,
        ];

        let outbound = 123u128.to_ne_bytes();
        let inbound = 234u128.to_ne_bytes();

        let alice = AeadTransport::<Neko128>::new(
            shared_secret.into(),
            outbound.into(),
            inbound.into(),
            Role::Sender,
        );
        let bob = AeadTransport::<Neko128>::new(
            shared_secret.into(),
            outbound.into(),
            inbound.into(),
            Role::Receiver,
        );

        let (mut alice_send, mut alice_recv) = alice.split();
        let (mut bob_send, mut bob_recv) = bob.split();

        let orig = b"Test Message, Please ignore. AAAAAAAAAAAAAAAAAAaaaaaaaaaaa | B .";

        let ad = b"random";

        let mut msg = orig.to_vec();

        // a -> b
        alice_send.encrypt(&mut msg, ad)?;

        assert_ne!(orig.as_slice(), msg.as_slice());
        let ct1 = msg.clone();

        bob_recv.decrypt(&mut msg, ad)?;

        // a -> b
        alice_send.encrypt(&mut msg, b"")?;

        assert_ne!(msg.as_slice(), ct1.as_slice());
        let ct2 = msg.clone();

        bob_recv.decrypt(&mut msg, b"")?;

        // b -> a
        bob_send.encrypt(&mut msg, ad)?;

        // None of the ciphertexts should match each other
        assert_ne!(msg.as_slice(), ct1.as_slice());
        assert_ne!(msg.as_slice(), ct2.as_slice());
        assert_ne!(ct1.as_slice(), ct2.as_slice());

        alice_recv.decrypt(&mut msg, ad)?;

        assert_eq!(orig.as_slice(), msg.as_slice());

        // Counters are tracked from sender to receiver
        assert_eq!(alice_send.counter, bob_recv.counter);
        assert_eq!(bob_send.counter, alice_recv.counter);

        // Counters are not linked on the same side
        assert_ne!(alice_send.counter, alice_recv.counter);
        assert_ne!(bob_send.counter, bob_recv.counter);

        Ok(())
    }

    #[test]
    fn two_way_transport_streaming_works() -> aead::Result<()> {
        let shared_secret = [
            0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d,
            0x8e, 0x8f, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9a, 0x9b,
            0x9c, 0x9d, 0x9e, 0x9f,
        ];

        let outbound = 123u128.to_ne_bytes();
        let inbound = 234u128.to_ne_bytes();

        let alice = AeadTransport::<Neko128>::new(
            shared_secret.into(),
            outbound.into(),
            inbound.into(),
            Role::Sender,
        );
        let bob = AeadTransport::<Neko128>::new(
            shared_secret.into(),
            outbound.into(),
            inbound.into(),
            Role::Receiver,
        );

        let (mut alice_send, mut alice_recv) = alice.split();
        let (mut bob_send, mut bob_recv) = bob.split();

        let mut alice_send_stream = alice_send.as_stream(b"strem")?;
        let mut bob_recv_stream = bob_recv.as_stream(b"strem")?;

        let mut message = b"OnePacket.TwoPacket.".to_vec();
        let orig = message.clone();

        alice_send_stream.produce(&mut message[..10]);
        alice_send_stream.produce(&mut message[10..]);
        let tag = alice_send_stream.finalize();

        assert_ne!(&message, &orig);
        let first_encrypted_msg = message.clone();

        bob_recv_stream.consume(&mut message[..10]);
        bob_recv_stream.consume(&mut message[10..]);
        bob_recv_stream.finalize(&tag)?;

        assert_eq!(&message, &orig);

        let mut bob_send_stream = bob_send.as_stream(b"strem")?;
        let mut alice_recv_stream = alice_recv.as_stream(b"strem")?;

        bob_send_stream.produce(&mut message[..10]);
        bob_send_stream.produce(&mut message[10..]);
        let tag = bob_send_stream.finalize();

        assert_ne!(&message, &orig);
        assert_ne!(&message, &first_encrypted_msg);

        alice_recv_stream.consume(&mut message[..10]);
        alice_recv_stream.consume(&mut message[10..]);
        alice_recv_stream.finalize(&tag)?;

        assert_eq!(&message, &orig);

        Ok(())
    }

    #[test]
    fn two_way_transport_fails_with_security_level_mismatch() {
        let shared_secret = [
            0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d,
            0x8e, 0x8f, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9a, 0x9b,
            0x9c, 0x9d, 0x9e, 0x9f,
        ];

        let outbound = 123u128.to_ne_bytes();
        let inbound = 234u128.to_ne_bytes();

        let alice = AeadTransport::<Neko128>::new(
            shared_secret.into(),
            outbound.into(),
            inbound.into(),
            Role::Sender,
        );
        let bob = AeadTransport::<Neko256>::new(
            shared_secret.into(),
            outbound.into(),
            inbound.into(),
            Role::Receiver,
        );

        let (mut alice_send, mut _alice_recv) = alice.split();
        let (mut _bob_send, mut bob_recv) = bob.split();

        let orig = b"Test Message, Please ignore.";

        let ad = b"random";

        let mut msg = orig.to_vec();

        // a -> b
        alice_send.encrypt(&mut msg, ad).unwrap();

        assert_ne!(orig.as_slice(), msg.as_slice());

        assert!(bob_recv.decrypt(&mut msg, ad).is_err());
    }
}