// This seems to be an absolute nightmare.  GenericArray sucks
// but I can't seem to figure out how to pull it out of this
// stack of software
use aes::Aes128;
#[allow(deprecated)]
use aes::cipher::{BlockDecrypt, BlockEncrypt, generic_array::GenericArray};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use curve25519_dalek::MontgomeryPoint;
use ed25519_dalek::{VerifyingKey, hazmat::ExpandedSecretKey};
use hmac::{Hmac, Mac};
use sha2::Sha256;

use crate::string_helper::NameString;

type HmacSha256 = Hmac<Sha256>;

pub trait Keystore {
    fn decrypt_and_id_p2p(
        &self,
        source: u8,
        _dest: u8,
        mac: u16,
        data: &Bytes,
    ) -> Option<(Bytes, u32, u32)>;

    fn decrypt_and_id_group(
        &self,
        group_hash_prefix: u8,
        mac: u16,
        data: &Bytes,
    ) -> Option<(Bytes, Option<NameString>)>;

    fn decrypt_anon(
        &self,
        dest: u8,
        pub_key: &PublicKey,
        mac: u16,
        data: &Bytes,
    ) -> Option<(Bytes, u32)>;
}

#[derive(Debug, PartialEq)]
pub enum MeshcoreCryptoError {
    KeyLengthError,
    TryFromSliceError,
    HexDecodeError,
    KeyCreationError,
}

impl core::error::Error for MeshcoreCryptoError {}

impl core::fmt::Display for MeshcoreCryptoError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            MeshcoreCryptoError::KeyLengthError => f.write_str("Key Length Error"),
            MeshcoreCryptoError::TryFromSliceError => f.write_str("Try From Slice Error"),
            MeshcoreCryptoError::HexDecodeError => f.write_str("Hex Decode Error"),
            MeshcoreCryptoError::KeyCreationError => f.write_str("Key Creation Error"),
        }
    }
}

#[derive(PartialEq)]
pub struct PrivateKey(ExpandedSecretKey);

impl core::fmt::Debug for PrivateKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_tuple("PrivateKey")
            .field(&hex::encode(self.0.scalar.as_bytes()))
            .finish()
    }
}

impl Clone for PrivateKey {
    fn clone(&self) -> Self {
        Self(ExpandedSecretKey {
            scalar: self.0.scalar,
            hash_prefix: self.0.hash_prefix,
        })
    }
}

impl Default for PrivateKey {
    fn default() -> Self {
        // To make a key whole-cloth, we need to start with a SigningKey made
        // using a good RNG.  Then we can use that to make an ExpandedSecretKey.
        use ed25519_dalek::SigningKey;
        use rand::rngs::OsRng;

        // This seems like the same sequence of steps that's used with ed25519 itself.
        // I have to copy-pasta it because I don't see a way to do it directly with
        // thei API.  In the docs the give this example for creating a signing key
        let mut csprng = OsRng;
        let signing_key: SigningKey = SigningKey::generate(&mut csprng);
        // Then, there are only a few constructors for making an ExpandedSecretKey.
        // Meshcore uses this kind of key, so it's what we need in this application,
        // but it's an uncommon formulation.
        let esk = ExpandedSecretKey::from(&signing_key.to_bytes());
        Self(esk)
    }
}

impl PrivateKey {
    pub fn hash_prefix(&self) -> u32 {
        // The has prefix is the beginning of the public key of the secret
        let public_key = PublicKey::from(self);
        public_key.hash_prefix()
    }
}

#[derive(PartialEq, Clone)]
pub struct PublicKey(ed25519_dalek::VerifyingKey);

impl core::fmt::Debug for PublicKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_tuple("PublicKey")
            .field(&hex::encode(self.0.as_bytes()))
            .finish()
    }
}

impl PublicKey {
    pub fn hash_prefix(&self) -> u32 {
        let mut bytes = Bytes::copy_from_slice(self.0.as_bytes());

        bytes.get_u32()
    }
}

#[derive(Clone, Eq, Hash)]
pub struct SharedSecret(MontgomeryPoint);

impl PartialEq for SharedSecret {
    fn eq(&self, other: &Self) -> bool {
        self.0.as_bytes() == other.0.as_bytes()
    }

    fn ne(&self, other: &Self) -> bool {
        !self.eq(other)
    }
}

impl core::fmt::Debug for SharedSecret {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_tuple("SharedSecret")
            .field(&hex::encode(self.0.as_bytes()))
            .finish()
    }
}

impl core::str::FromStr for SharedSecret {
    type Err = MeshcoreCryptoError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut array = [0_u8; 32];

        // The provided group secrets are only 16 bytes,
        // but they're zero-paded to be 32.  So, we're
        // going to get the hex from the string and copy it in.
        if hex::decode_to_slice(s, &mut array[0..16]).is_err() {
            return Err(MeshcoreCryptoError::TryFromSliceError);
        } else {
            Ok(SharedSecret(MontgomeryPoint(array)))
        }
    }
}

impl TryFrom<Bytes> for SharedSecret {
    type Error = MeshcoreCryptoError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if value.len() != 32 {
            Err(MeshcoreCryptoError::KeyLengthError)
        } else {
            let mut value = value;
            let mut retval = Self(MontgomeryPoint([0_u8; 32]));
            value.copy_to_slice(&mut retval.0.0);
            Ok(retval)
        }
    }
}

// This is kinda meaningless because a shared secret only works
// when it's connected to a key pair, but it needs to exist for
// Arrayvec.
impl Default for SharedSecret {
    fn default() -> Self {
        Self(curve25519_dalek::MontgomeryPoint([0u8; 32]))
    }
}

impl SharedSecret {
    fn get_key(&self) -> &[u8; 16] {
        // Safety: The size of the slice ensures that this will never be wrong.
        (&self.0.as_bytes()[0..16]).try_into().unwrap()
    }

    pub fn new_from_group_secret(bytes: Bytes) -> Self {
        // The group secret is 16-bytes of key with the last 16-bytes set to zero.
        let mut group_secret = BytesMut::from(bytes);
        group_secret.reserve(16);
        group_secret.put(&[0_u8; 16][..]);

        let mut slice = [0_u8; 32];
        group_secret.copy_to_slice(&mut slice);

        SharedSecret(MontgomeryPoint(slice))
    }

    pub fn get_hmac(&self, ciphertext: &Bytes) -> u16 {
        // At this point, we're sure that the key is an appropriate size for
        // the Hmac.  So, I don't think we should complicate the API with making
        // this failable.
        let mut mac =
            HmacSha256::new_from_slice(self.0.as_bytes()).expect("Programming error in hmac");
        mac.update(&ciphertext);
        let result = mac.finalize();
        let mut bytes = Bytes::copy_from_slice(&result.into_bytes());
        bytes.get_u16()
    }

    pub fn decrypt(&self, ciphertext: &Bytes) -> Bytes {
        use aes::cipher::KeyInit;

        if let Ok(aes) = Aes128::new_from_slice(self.get_key()) {
            let mut text = BytesMut::from(ciphertext.clone());

            // The decryption function works on a 16-byte block of data.
            // we need to break the input ciphertext into blocks of this size
            // and work with them individually.  Also, any "short" blocks have
            // to be zero-padded.
            // Copy the original length so we can truncate back down
            let chunk_size = 16;
            let length = text.len();
            let remainder = length % chunk_size;

            // Pad the bytes to it's an even multiple of the chunk size
            if remainder != 0 {
                let padding = chunk_size - remainder;
                text.reserve(padding);
                text.put_bytes(0, padding);
            }

            let chunks = text.chunks_exact_mut(chunk_size);
            for chunk in chunks {
                #[allow(deprecated)]
                let mut block = *GenericArray::from_slice(&chunk);
                aes.decrypt_block(&mut block);
                chunk.copy_from_slice(&block);
            }

            // Drop the padding
            text.truncate(length);

            Bytes::from(text)
        } else {
            assert!(false, "Key length error in decrypt");
            Bytes::new()
        }
    }

    pub fn encrypt(&self, plaintext: Bytes) -> Bytes {
        use aes::cipher::KeyInit;

        // Safety: The get_key function is guarenteed to produce a key
        // the right size.  That's the only reason this method might
        // fail.
        let aes = Aes128::new_from_slice(self.get_key()).unwrap();
        let mut text = BytesMut::from(plaintext);

        // The decryption function works on a 16-byte block of data.
        // we need to break the input ciphertext into blocks of this size
        // and work with them individually.  Also, any "short" blocks have
        // to be zero-padded.
        // Copy the original length so we can truncate back down
        let chunk_size = 16;
        let length = text.len();
        let remainder = length % chunk_size;

        // Pad the bytes to it's an even multiple of the chunk size
        if remainder != 0 {
            let padding = chunk_size - remainder;
            text.reserve(padding);
            text.put_bytes(0, padding);
        }

        let chunks = text.chunks_exact_mut(chunk_size);
        for chunk in chunks {
            #[allow(deprecated)]
            let mut block = *GenericArray::from_slice(&chunk);
            aes.encrypt_block(&mut block);
            chunk.copy_from_slice(&block);
        }

        Bytes::from(text)
    }

    pub fn hash_prefix(&self) -> u8 {
        self.0.0[0]
    }
}

// This is just for creating placeholders
impl Default for PublicKey {
    fn default() -> Self {
        PublicKey(VerifyingKey::from_bytes(&[0_u8; 32]).unwrap())
    }
}

impl TryFrom<&[u8]> for PublicKey {
    type Error = MeshcoreCryptoError;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        let bytes = Bytes::copy_from_slice(value);
        Self::try_from(bytes)
    }
}

impl TryFrom<Bytes> for PublicKey {
    type Error = MeshcoreCryptoError;

    fn try_from(mut value: Bytes) -> Result<Self, Self::Error> {
        if value.len() < 32 {
            return Err(MeshcoreCryptoError::KeyLengthError);
        }

        let mut slice = [0_u8; 32];

        if value.try_copy_to_slice(&mut slice).is_err() {
            return Err(MeshcoreCryptoError::KeyLengthError);
        }

        if let Ok(key) = VerifyingKey::from_bytes(&slice) {
            Ok(PublicKey(key))
        } else {
            Err(MeshcoreCryptoError::KeyCreationError)
        }
    }
}

impl core::str::FromStr for PublicKey {
    type Err = MeshcoreCryptoError;

    fn from_str(hex_str: &str) -> Result<Self, Self::Err> {
        if let Ok(hex) = hex::decode(hex_str) {
            if let Ok(slice) = TryInto::<[u8; 32]>::try_into(hex) {
                if let Ok(key) = VerifyingKey::from_bytes(&slice) {
                    Ok(PublicKey(key))
                } else {
                    Err(MeshcoreCryptoError::KeyCreationError)
                }
            } else {
                Err(MeshcoreCryptoError::KeyLengthError)
            }
        } else {
            Err(MeshcoreCryptoError::HexDecodeError)
        }
    }
}

impl core::str::FromStr for PrivateKey {
    type Err = MeshcoreCryptoError;

    fn from_str(hex_str: &str) -> Result<Self, Self::Err> {
        if let Ok(hex) = hex::decode(hex_str) {
            if let Ok(bytes) = TryInto::<[u8; 64]>::try_into(hex) {
                Ok(PrivateKey(ExpandedSecretKey::from_bytes(&bytes)))
            } else {
                Err(MeshcoreCryptoError::TryFromSliceError)
            }
        } else {
            Err(MeshcoreCryptoError::HexDecodeError)
        }
    }
}

impl From<&PrivateKey> for PublicKey {
    fn from(key: &PrivateKey) -> Self {
        let key = (key.0.scalar * curve25519_dalek::constants::ED25519_BASEPOINT_POINT).compress();
        PublicKey(VerifyingKey::from_bytes(key.as_bytes()).unwrap())
    }
}

impl PrivateKey {
    pub fn create_secret(&self, other: &PublicKey) -> SharedSecret {
        SharedSecret(self.0.scalar * other.0.to_montgomery())
    }
}

impl SharedSecret {
    pub fn mac_then_decrypt(&self, mac: u16, data: &Bytes) -> Option<Bytes> {
        // Get the MAC of the message and key to check vailidity
        let our_mac = self.get_hmac(&data);
        if our_mac != mac {
            return None;
        }

        // Attempt to decrypt the packet itself
        Some(self.decrypt(&data))
    }

    pub fn encrypt_then_mac(&self, data: Bytes) -> (u16, Bytes) {
        let ciphertext = self.encrypt(data);
        let mac = self.get_hmac(&ciphertext);
        (mac, ciphertext)
    }
}

// Tests for std operations
#[cfg(test)]
mod tests {
    use super::*;
    use core::str::FromStr;
    use hex::{decode, decode_to_slice, encode};

    #[test]
    fn public_key() {
        let mut slice = [0_u8; 32];
        decode_to_slice(
            "12349bdc1f76a0c12149bb15f791dbe42fde02c209b04a85c6f512990c8cedec",
            &mut slice,
        )
        .unwrap();
        let public_key =
            PublicKey::from_str("12349bdc1f76a0c12149bb15f791dbe42fde02c209b04a85c6f512990c8cedec");
        assert_eq!(
            Ok(PublicKey(VerifyingKey::from_bytes(&slice).unwrap())),
            public_key
        );
    }

    #[test]
    fn private_key() {
        let private_key = PrivateKey::from_str("38DAA98490B7284697C7ADA6175FD1F8DAD12032AD7ABAE625B7EAD8FEC6444CA281C3370B97155D9C8CECD89A929FDDE0FBF3A9D5C92A1B3C24D711934CD69D").unwrap();
        let public_key = PublicKey::from(&private_key);
        println!("Public key: {:#?}", public_key);
        assert_eq!(
            PublicKey::from_str("12349bdc1f76a0c12149bb15f791dbe42fde02c209b04a85c6f512990c8cedec")
                .unwrap(),
            public_key
        );
    }

    #[test]
    fn shared_secret() {
        // We'll make a public/private pair for alice and bob and make sure the shared secret is the same and expected
        let alice_private = PrivateKey::from_str("4885CF25975EA09742EF76DA587D0957E74EE02AAA34A001458E207E63CF7E6C4940C8C42C335862C71CC2F139633057D1FEE5687B172B27E1E0302A1D480E08").unwrap();
        let bob_private = PrivateKey::from_str("38DAA98490B7284697C7ADA6175FD1F8DAD12032AD7ABAE625B7EAD8FEC6444CA281C3370B97155D9C8CECD89A929FDDE0FBF3A9D5C92A1B3C24D711934CD69D").unwrap();

        let alice_public = PublicKey::from(&alice_private);
        let bob_public = PublicKey::from(&bob_private);

        assert_eq!(
            alice_public.0.as_bytes().to_vec(),
            decode("34569df1f9661916901669666fb8025eccb9ddb0499cddad4c164fec219c8b8f").unwrap()
        );
        assert_eq!(
            bob_public.0.as_bytes().to_vec(),
            decode("12349bdc1f76a0c12149bb15f791dbe42fde02c209b04a85c6f512990c8cedec").unwrap()
        );

        println!("Alice's public key: {}", encode(&alice_public.0.to_bytes()));
        println!("Bob's   public key: {}", encode(&bob_public.0.to_bytes()));

        let left_secret = alice_private.create_secret(&bob_public);
        let right_secret = bob_private.create_secret(&alice_public);

        assert_eq!(
            left_secret.0.as_bytes().to_vec(),
            decode("eb7a365363bd8548ee2b54b9234247be5e42e96be9625adcdf3a55b6c1d04850").unwrap()
        );

        println!("Left shared secret:  {}", encode(&left_secret.0.as_bytes()));
        println!(
            "Right shared secret: {}",
            encode(&right_secret.0.as_bytes())
        );

        assert_eq!(left_secret, right_secret);
    }

    #[test]
    fn hmac() {
        // Test using a group secret
        let group_secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(
            &decode("8b3387e9c5cdea6ac9e5edbaa115cd72").unwrap(),
        ));
        let test_group_secret = SharedSecret::from_str("8b3387e9c5cdea6ac9e5edbaa115cd72").unwrap();
        assert_eq!(group_secret, test_group_secret);

        // Test using the secret and ciphertext to make a MAC and ensure it matches an example for a group secret
        let sample_data = Bytes::copy_from_slice(
            &decode("354D619BAE9590E4D177DB7EEAF982F5BDCF78005D75157D9535FA90178F785D").unwrap(),
        );
        let mac = group_secret.get_hmac(&sample_data);
        assert_eq!(0xC3C1, mac);
    }

    #[test]
    fn decrypt() {
        let ciphertext = Bytes::copy_from_slice(
            &decode("354D619BAE9590E4D177DB7EEAF982F5BDCF78005D75157D9535FA90178F785D").unwrap(),
        );
        let secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(
            &decode("8b3387e9c5cdea6ac9e5edbaa115cd72").unwrap(),
        ));
        let cleartext = secret.decrypt(&ciphertext);
        println!("Cleartext: {}", encode(&cleartext));
    }

    #[test]
    fn decrypt_online_example() {
        let ciphertext =
            Bytes::copy_from_slice(&decode("9A1FD57EDFE7E4369F9FD9420C48FFAD").unwrap());
        let secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(
            &decode("949E911CA6A6196275FF319B28C3A143").unwrap(),
        ));
        let cleartext = secret.decrypt(&ciphertext);
        let vec = cleartext.to_vec();
        let string = String::from_utf8_lossy(&vec);
        assert_eq!("Hello my world!!", string);
    }

    #[test]
    fn encrypt_online_example() {
        let plaintext = Bytes::copy_from_slice("Meshcore!".as_bytes());
        let secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(
            &decode("44A6F78DAD2E54D73A32CDE3ECAA9E75").unwrap(),
        ));
        let ciphertext = secret.encrypt(plaintext);
        assert_eq!(
            ciphertext,
            decode("62374852B6A11405A081F87356C88861").unwrap()
        );
    }

    #[test]
    fn chunks() {
        // Test my understanding of how Byte's chunk_exact methods work for use in (en|de)cryption
        let array = decode("0102030405060708090A0B0C0D0E0F").unwrap();
        let mut text = BytesMut::from(array.as_slice());

        // Copy the original length so we can truncate back down
        let chunk_size = 4;
        let length = text.len();

        // Pad the bytes to it's an even multiple of the chunk size
        let padding = chunk_size - (length % chunk_size);
        text.reserve(padding);
        text.put_bytes(0, padding);

        let chunks = text.chunks_exact_mut(4);
        {
            for chunk in chunks {
                for item in chunk {
                    *item = *item << 4;
                }
            }
        }

        // Drop the padding
        text.truncate(length);

        println!("Result: {:#?}", encode(&text));

        assert_eq!(text, decode("102030405060708090A0B0C0D0E0F0").unwrap());
    }

    // Example from the crate we're using
    #[test]
    fn aes_test() {
        use aes::Aes128;
        #[allow(deprecated)]
        use aes::cipher::{BlockDecrypt, BlockEncrypt, KeyInit};

        // Initialize cipher
        let key: [u8; 16] = decode("0A1BB8C05063D9941F0F1019D001B743")
            .unwrap()
            .try_into()
            .unwrap();
        let cipher = Aes128::new(&key.into());

        let message: [u8; 16] = decode("7D69AB072E09AF74EBA47EB95BF00AE3")
            .unwrap()
            .try_into()
            .unwrap();
        let message_copy = message.clone();

        // Encrypt block in-place
        println!("Before: {}", encode(&message));

        cipher.encrypt_block(&mut message.into());

        println!("Crypted: {}", encode(&message));

        cipher.decrypt_block(&mut message.into());

        println!("Decrypted: {}", encode(&message));

        assert_eq!(message, message_copy);
    }

    #[test]
    fn shared_secrets_example() {
        use x25519_dalek::{PublicKey, StaticSecret};

        let a_s: [u8; 32] =
            decode("58f5052c13275c8a3f4863a082555fed7ea08b9dec2eb00dd86f6b5412174458")
                .unwrap()
                .try_into()
                .unwrap();
        let b_s: [u8; 32] =
            decode("586c4cc29635af5865abe4f231bafbd373969725493c07271e02c7fec8ff3b5f")
                .unwrap()
                .try_into()
                .unwrap();
        let alice_secret = StaticSecret::from(a_s);
        let alice_public = PublicKey::from(&alice_secret);

        let bob_secret = StaticSecret::from(b_s);
        let bob_public = PublicKey::from(&bob_secret);

        let alice_shared_secret = alice_secret.diffie_hellman(&bob_public);
        let bob_shared_secret = bob_secret.diffie_hellman(&alice_public);

        assert_eq!(alice_shared_secret.as_bytes(), bob_shared_secret.as_bytes());
    }

    #[test]
    fn mac_then_decrypt() {
        let group_secret = SharedSecret::new_from_group_secret(Bytes::copy_from_slice(
            &decode("8b3387e9c5cdea6ac9e5edbaa115cd72").unwrap(),
        ));
        let sample_data = Bytes::copy_from_slice(
            &decode("354D619BAE9590E4D177DB7EEAF982F5BDCF78005D75157D9535FA90178F785D").unwrap(),
        );
        let mac = 0xC3C1;

        let cleartext = group_secret.mac_then_decrypt(mac, &sample_data).unwrap();
        assert_eq!(
            cleartext,
            decode("3757d06800f09f8cb220547265653a20e29881efb88f00000000000000000000").unwrap()
        );
    }

    #[test]
    fn test_error_display() {
        assert_eq!(
            format!("{}", MeshcoreCryptoError::KeyLengthError),
            "Key Length Error"
        );
        assert_eq!(
            format!("{}", MeshcoreCryptoError::TryFromSliceError),
            "Try From Slice Error"
        );
        assert_eq!(
            format!("{}", MeshcoreCryptoError::HexDecodeError),
            "Hex Decode Error"
        );
        assert_eq!(
            format!("{}", MeshcoreCryptoError::KeyCreationError),
            "Key Creation Error"
        );
    }

    #[test]
    fn test_crypto_round_trip() {
        let alice_private = PrivateKey::default();
        let alice_public = PublicKey::from(&alice_private);

        let bob_private = PrivateKey::default();
        let bob_public = PublicKey::from(&bob_private);

        let cleartext = "Hi Bob.  This is alice.  How are you?";
        let alice_secret = alice_private.create_secret(&bob_public);

        let (mac, ciphertext) =
            alice_secret.encrypt_then_mac(Bytes::copy_from_slice(cleartext.as_bytes()));

        let bob_secret = bob_private.create_secret(&alice_public);
        let decrypted_data = bob_secret.mac_then_decrypt(mac, &ciphertext);

        if let Some(decrypted_data) = decrypted_data {
            // Due to padding for the block cipher, the decrypted string may have up to
            // 31 null characters at the end.  Given that we want to compare the starting
            // and ending, we have to strip these back out.  We can't do this in the decrypt
            // function because we don't actually know how long the contents of the is
            // suppose to be.
            let raw_message = String::from_utf8_lossy(&decrypted_data).to_owned();
            let lhs_string: &str = raw_message
                .splitn(2, "\0")
                .collect::<Vec<&str>>()
                .first()
                .unwrap();

            assert_eq!(lhs_string, cleartext);
        } else {
            assert!(false, "Unable to decrypt");
        }
    }
}