// SPDX-License-Identifier: GPL-3.0-only

use chacha20poly1305::{KeyInit, XChaCha20Poly1305, XNonce, aead::AeadInPlace};
use std::io::Write;

use crate::error::*;
use crate::header::{AlgId, Header, KdfParams, NONCE_PREFIX_LEN, TAG_LEN};
use crate::reader::{AheadReader, ReadInfoChunk};
use crate::secrets::SecretBytes32;
use crate::utils::*;

/// XChaCha20Poly1305 nonce: 24 bytes total. STREAM splits the trailing 5 bytes
/// into a 4-byte big-endian counter and a 1-byte "last block" flag.
const NONCE_LEN: usize = 24;
const COUNTER_LEN: usize = 4;
const _: () = assert!(NONCE_PREFIX_LEN + COUNTER_LEN + 1 == NONCE_LEN);

fn make_nonce(prefix: &[u8; NONCE_PREFIX_LEN], counter: u32, last: bool) -> XNonce {
    let mut n = [0u8; NONCE_LEN];
    n[..NONCE_PREFIX_LEN].copy_from_slice(prefix);
    n[NONCE_PREFIX_LEN..NONCE_PREFIX_LEN + COUNTER_LEN].copy_from_slice(&counter.to_be_bytes());
    n[NONCE_LEN - 1] = u8::from(last);
    XNonce::from(n)
}

/// Derive (or unwrap) the 32-byte AEAD key from KDF parameters and an optional passphrase.
/// For `KdfParams::Raw`, `raw_key` must be supplied.
/// For `KdfParams::Argon2id`, `passphrase` must be supplied.
pub fn derive_key(
    kdf: &KdfParams,
    raw_key: Option<&[u8; 32]>,
    passphrase: Option<&[u8]>,
) -> Result<SecretBytes32, FcryError> {
    let mut out = SecretBytes32::zeroed();
    match kdf {
        KdfParams::Raw => {
            let raw =
                raw_key.ok_or_else(|| FcryError::Format("raw kdf requires --raw-key".into()))?;
            out.as_mut_array().copy_from_slice(raw);
        }
        KdfParams::Argon2id {
            salt,
            m_cost,
            t_cost,
            p_cost,
        } => {
            let pw = passphrase
                .ok_or_else(|| FcryError::Format("argon2id kdf requires a passphrase".into()))?;
            let params = argon2::Params::new(*m_cost, *t_cost, *p_cost, Some(32))?;
            let argon =
                argon2::Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
            argon.hash_password_into(pw, salt, out.as_mut_array())?;
        }
    }
    Ok(out)
}

pub fn encrypt<S: AsRef<str>>(
    input_file: Option<S>,
    output_file: Option<S>,
    key: &SecretBytes32,
    chunk_size: u32,
    kdf: KdfParams,
) -> Result<(), FcryError> {
    let chunk_sz = chunk_size as usize;
    let mut f_plain = AheadReader::from(open_input(input_file)?, chunk_sz);
    let mut f_encrypted = OutSink::open(output_file)?;

    let mut nonce_prefix = [0u8; NONCE_PREFIX_LEN];
    getrandom::fill(&mut nonce_prefix)?;

    let header = Header {
        alg: AlgId::XChaCha20Poly1305,
        flags: 0,
        chunk_size,
        kdf,
        nonce_prefix,
    };
    let aad = header.encode();
    f_encrypted.write_all(&aad)?;

    let aead = XChaCha20Poly1305::new(key.as_array().into());

    let mut buf = vec![0u8; chunk_sz];
    let mut counter: u32 = 0;

    loop {
        match f_plain.read_ahead(&mut buf)? {
            ReadInfoChunk::Normal(_) => {
                let nonce = make_nonce(&nonce_prefix, counter, false);
                aead.encrypt_in_place(&nonce, &aad, &mut buf)?;
                f_encrypted.write_all(&buf)?;
                buf.truncate(chunk_sz);
                counter = counter.checked_add(1).ok_or_else(|| {
                    FcryError::Format("STREAM counter overflow (input too large)".into())
                })?;
            }
            ReadInfoChunk::Last(n) => {
                buf.truncate(n);
                let nonce = make_nonce(&nonce_prefix, counter, true);
                aead.encrypt_in_place(&nonce, &aad, &mut buf)?;
                f_encrypted.write_all(&buf)?;
                break;
            }
            ReadInfoChunk::Empty => {
                // Empty plaintext: still emit a final "last" tag so the decryptor
                // authenticates the (empty) stream rather than silently producing nothing.
                buf.clear();
                let nonce = make_nonce(&nonce_prefix, counter, true);
                aead.encrypt_in_place(&nonce, &aad, &mut buf)?;
                f_encrypted.write_all(&buf)?;
                break;
            }
        }
    }

    f_encrypted.commit()?;
    Ok(())
}

pub fn decrypt<S: AsRef<str>>(
    input_file: Option<S>,
    output_file: Option<S>,
    raw_key: Option<&[u8; 32]>,
    passphrase: Option<&[u8]>,
) -> Result<(), FcryError> {
    let mut reader = open_input(input_file)?;
    let header = Header::read(&mut reader)?;
    let aad = header.encode();

    let key = derive_key(&header.kdf, raw_key, passphrase)?;

    let chunk_sz = header.chunk_size as usize;
    let cipher_chunk = chunk_sz + TAG_LEN;

    let mut f_encrypted = AheadReader::from(reader, cipher_chunk);
    let mut f_plain = OutSink::open(output_file)?;

    let aead = XChaCha20Poly1305::new(key.as_array().into());

    let mut buf = vec![0u8; cipher_chunk];
    let mut counter: u32 = 0;

    loop {
        match f_encrypted.read_ahead(&mut buf)? {
            ReadInfoChunk::Normal(_) => {
                let nonce = make_nonce(&header.nonce_prefix, counter, false);
                aead.decrypt_in_place(&nonce, &aad, &mut buf)?;
                f_plain.write_all(&buf)?;
                buf.resize(cipher_chunk, 0);
                counter = counter
                    .checked_add(1)
                    .ok_or_else(|| FcryError::Format("STREAM counter overflow".into()))?;
            }
            ReadInfoChunk::Last(n) => {
                buf.truncate(n);
                let nonce = make_nonce(&header.nonce_prefix, counter, true);
                aead.decrypt_in_place(&nonce, &aad, &mut buf)?;
                f_plain.write_all(&buf)?;
                break;
            }
            ReadInfoChunk::Empty => {
                return Err(FcryError::Format(
                    "truncated ciphertext: missing final chunk".into(),
                ));
            }
        }
    }

    f_plain.commit()?;
    Ok(())
}