use std::io::ErrorKind;

use aes_gcm::aead::{Aead, OsRng};
use aes_gcm::{Aes256Gcm, Key, KeyInit, Nonce};
use rand::Rng;
use serde::de::DeserializeOwned;
use serde::Serialize;

use crate::utils::err::Res;

const NONCE_LEN: usize = 12;

pub struct CryptoWrapper {
    key: Key<Aes256Gcm>,
}

impl CryptoWrapper {
    /// Generate a new memory wrapper
    pub fn new_random() -> Self {
        Self {
            key: Aes256Gcm::generate_key(&mut OsRng),
        }
    }

    /// Encrypt some data
    pub fn encrypt<T: Serialize + DeserializeOwned>(&self, data: &T) -> Res<String> {
        let aes_key = Aes256Gcm::new(&self.key);
        let nonce_bytes = rand::thread_rng().gen::<[u8; NONCE_LEN]>();

        let serialized_data = bincode::serialize(data)?;

        let mut enc = aes_key
            .encrypt(Nonce::from_slice(&nonce_bytes), serialized_data.as_slice())
            .unwrap();
        enc.extend_from_slice(&nonce_bytes);

        Ok(base64::encode(enc))
    }

    /// Decrypt some data previously encrypted using the [`CryptoWrapper::encrypt`] method
    pub fn decrypt<T: DeserializeOwned>(&self, input: &str) -> Res<T> {
        let bytes = base64::decode(input)?;

        if bytes.len() < NONCE_LEN {
            return Err(Box::new(std::io::Error::new(
                ErrorKind::Other,
                "Input string is smaller than nonce!",
            )));
        }

        let (enc, nonce) = bytes.split_at(bytes.len() - NONCE_LEN);
        assert_eq!(nonce.len(), NONCE_LEN);

        let aes_key = Aes256Gcm::new(&self.key);

        let dec = match aes_key.decrypt(Nonce::from_slice(nonce), enc) {
            Ok(d) => d,
            Err(e) => {
                log::error!("Failed to decrypt wrapped data! {:#?}", e);
                return Err(Box::new(std::io::Error::new(
                    ErrorKind::Other,
                    "Failed to decrypt wrapped data!",
                )));
            }
        };

        Ok(bincode::deserialize(&dec)?)
    }
}

#[cfg(test)]
mod test {
    use crate::data::crypto_wrapper::CryptoWrapper;

    #[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)]
    struct Message(String);

    #[test]
    fn encrypt_and_decrypt() {
        let wrapper = CryptoWrapper::new_random();
        let msg = Message("Pierre was here".to_string());
        let enc = wrapper.encrypt(&msg).unwrap();
        let dec: Message = wrapper.decrypt(&enc).unwrap();

        assert_eq!(dec, msg)
    }

    #[test]
    fn encrypt_and_decrypt_invalid() {
        let wrapper_1 = CryptoWrapper::new_random();
        let wrapper_2 = CryptoWrapper::new_random();
        let msg = Message("Pierre was here".to_string());
        let enc = wrapper_1.encrypt(&msg).unwrap();
        wrapper_2.decrypt::<Message>(&enc).unwrap_err();
    }
}