use std::ptr::NonNull;

use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
use rand::{rng, Rng};

pub(super) const KEY_SIZE: usize = 32;
pub(super) const NONCE_SIZE: usize = 12;

/// The encryption performed here is AES-256-GCM. Any tampering with the key or the ciphertexts
/// will result in a decryption failure and panic. The key's memory contents are protected from
/// being swapped to disk via mlock.
pub(super) struct MemoryEncryptionKey(NonNull<[u8]>);

/// An encrypted memory blob that must be decrypted using the same key that it was encrypted with.
pub struct EncryptedMemory {
    nonce: [u8; NONCE_SIZE],
    ciphertext: Vec<u8>,
}

impl MemoryEncryptionKey {
    pub fn new() -> Self {
        let mut key = [0u8; KEY_SIZE];
        rng().fill(&mut key);
        MemoryEncryptionKey::from(&key)
    }

    /// Encrypts the given plaintext using the key.
    #[allow(unused)]
    pub(super) fn encrypt(&self, plaintext: &[u8]) -> EncryptedMemory {
        let cipher = Aes256Gcm::new_from_slice(self.as_ref()).expect("Could not create aes key");
        let mut nonce = [0u8; NONCE_SIZE];
        rng().fill(&mut nonce);
        let ciphertext = cipher
            .encrypt(&Nonce::from(nonce), plaintext)
            .expect("encryption should not fail");
        EncryptedMemory { nonce, ciphertext }
    }

    /// Decrypts the given encrypted memory using the key. A decryption failure will panic. This is
    /// okay because neither the keys nor ciphertexts should ever fail to decrypt, and doing so
    /// indicates that the process memory was tampered with.
    #[allow(unused)]
    pub(super) fn decrypt(&self, encrypted: &EncryptedMemory) -> Result<Vec<u8>, DecryptionError> {
        let cipher = Aes256Gcm::new_from_slice(self.as_ref()).expect("Could not create aes key");
        cipher
            .decrypt(&Nonce::from(encrypted.nonce), encrypted.ciphertext.as_ref())
            .map_err(|_| DecryptionError::CouldNotDecrypt)
    }
}

impl Drop for MemoryEncryptionKey {
    fn drop(&mut self) {
        unsafe {
            memsec::free(self.0);
        }
    }
}

impl From<&[u8; KEY_SIZE]> for MemoryEncryptionKey {
    fn from(value: &[u8; KEY_SIZE]) -> Self {
        let mut ptr: NonNull<[u8]> =
            unsafe { memsec::malloc_sized(KEY_SIZE).expect("malloc_sized should work") };
        unsafe {
            std::ptr::copy_nonoverlapping(value.as_ptr(), ptr.as_mut().as_mut_ptr(), KEY_SIZE);
        }
        MemoryEncryptionKey(ptr)
    }
}

impl AsRef<[u8]> for MemoryEncryptionKey {
    fn as_ref(&self) -> &[u8] {
        unsafe { self.0.as_ref() }
    }
}

#[derive(Debug)]
pub enum DecryptionError {
    CouldNotDecrypt,
}

impl std::error::Error for DecryptionError {}

impl std::fmt::Display for DecryptionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DecryptionError::CouldNotDecrypt => write!(f, "Could not decrypt"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_memory_encryption_key() {
        let key = MemoryEncryptionKey::new();
        let data = b"Hello, world!";
        let encrypted = key.encrypt(data);
        let decrypted = key.decrypt(&encrypted).unwrap();
        assert_eq!(data.as_ref(), decrypted.as_slice());
    }

    #[test]
    fn test_tampered_ciphertext_fails_to_decrypt() {
        let key = MemoryEncryptionKey::new();
        let mut encrypted = key.encrypt(b"Hello, world!");
        encrypted.ciphertext[0] ^= 0xff;
        let result = key.decrypt(&encrypted);
        assert!(matches!(result, Err(DecryptionError::CouldNotDecrypt)));
    }
}
