use std::io::{Read, Write};

use aes_gcm::{
    aead::{generic_array::GenericArray, Aead},
    Aes256Gcm, KeyInit,
};
use base64::Engine;
use pbkdf2::{password_hash::PasswordHasher, Pbkdf2};
use rand::{thread_rng, RngCore};

use crate::Error;

pub const SALT_LEN: usize = 24;
pub const NONCE_LEN: usize = 12;

pub fn random_bytes(length: usize) -> Vec<u8> {
    let mut rng = thread_rng();
    let mut data = vec![0u8; length];
    rng.fill_bytes(&mut data);
    return data;
}

pub fn derive_key(password: &[u8], salt: &[u8]) -> Result<Vec<u8>, Error> {
    let salt = base64::engine::general_purpose::STANDARD.encode(salt);

    let hash = Pbkdf2.hash_password(password, &salt)?.hash;
    let hash_bytes = hash.expect("Password hash empty").as_bytes().to_owned();
    return Ok(hash_bytes);
}

pub fn encrypt<W: Write>(mut writer: W, data: &[u8], key: &[u8]) -> Result<(), Error> {
    let cipher = Aes256Gcm::new(GenericArray::from_slice(key));

    let nonce = random_bytes(NONCE_LEN);
    let ciphertext = cipher.encrypt(GenericArray::from_slice(&nonce), data)?;

    writer.write_all(&nonce)?;
    writer.write_all(&ciphertext)?;

    Ok(())
}

pub fn decrypt<R: Read>(mut reader: R, key: &[u8]) -> Result<Vec<u8>, Error> {
    let mut nonce = [0u8; NONCE_LEN];
    reader.read_exact(&mut nonce)?;

    let mut ciphertext = Vec::with_capacity(1024);
    reader.read_to_end(&mut ciphertext)?;

    let cipher = Aes256Gcm::new(GenericArray::from_slice(key));
    let plaintext = cipher.decrypt(GenericArray::from_slice(&nonce), &ciphertext[..])?;
    Ok(plaintext)
}