// SPDX-License-Identifier: GPL-3.0-only
use chacha20poly1305::{
    aead::{stream, NewAead},
    ChaCha20Poly1305,
};
use rand::{rngs::OsRng, RngCore};

use crate::error::*;
use crate::utils::*;

const BUFSIZE: usize = 64 * 1024; // 64 KiB

fn new_random_nonce() -> Result<[u8; 7], FcryError> {
    let mut nonce = [0u8; 7];
    OsRng.try_fill_bytes(&mut nonce)?;
    Ok(nonce)
}

pub fn encrypt<S: AsRef<str>>(
    input_file: Option<S>,
    output_file: Option<S>,
    key: [u8; 32],
) -> Result<(), FcryError> {
    let mut f_plain = read_from_file_or_stdin(input_file);
    let mut f_encrypted = write_to_file_or_stdout(output_file);

    let nonce = new_random_nonce()?;
    f_encrypted.write_all(&nonce)?;

    let aead = ChaCha20Poly1305::new(&key.into());
    let mut stream_encryptor = stream::EncryptorBE32::from_aead(aead, &nonce.into());

    let mut buf = vec![0; BUFSIZE];
    let mut read_bytes;

    loop {
        read_bytes = f_plain.read(&mut buf)?;

        if read_bytes < BUFSIZE {
            break;
        }

        stream_encryptor.encrypt_next_in_place(&[], &mut buf)?;
        f_encrypted.write_all(&buf)?;
        buf.truncate(BUFSIZE);
    }

    buf.truncate(read_bytes);
    stream_encryptor.encrypt_last_in_place(&[], &mut buf)?;
    f_encrypted.write_all(&buf)?;

    Ok(())
}

pub fn decrypt<S: AsRef<str>>(
    input_file: Option<S>,
    output_file: Option<S>,
    key: [u8; 32],
) -> Result<(), FcryError> {
    let mut f_encrypted = read_from_file_or_stdin(input_file);
    let mut f_plain = write_to_file_or_stdout(output_file);

    let mut nonce = [0u8; 7];
    f_encrypted.read_exact(&mut nonce)?;

    let aead = ChaCha20Poly1305::new(&key.into());
    let mut stream_decryptor = stream::DecryptorBE32::from_aead(aead, &nonce.into());

    let mut buf = vec![0; BUFSIZE + 16];
    let mut read_bytes;

    loop {
        read_bytes = f_encrypted.read(&mut buf)?;

        if read_bytes < BUFSIZE + 16 {
            break;
        }

        stream_decryptor.decrypt_next_in_place(&[], &mut buf)?;
        f_plain.write_all(&buf)?;
        buf.resize(BUFSIZE + 16, 0);
    }

    buf.truncate(read_bytes);
    stream_decryptor.decrypt_last_in_place(&[], &mut buf)?;
    f_plain.write_all(&buf)?;

    Ok(())
}