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

use chacha20poly1305::{KeyInit, XChaCha20Poly1305, XNonce, aead::AeadInPlace};
use std::fs::File;
use std::io::{BufReader, Read, Seek, SeekFrom, Write};
use std::sync::Arc;

use crate::error::*;
use crate::header::{AlgId, FLAG_LENGTH_COMMITTED, Header, KdfParams, NONCE_PREFIX_LEN, TAG_LEN};
use crate::pipeline;
use crate::reader::{AheadReader, ReadInfoChunk};
use crate::secrets::{SecretBytes32, SecretVec};
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.
pub(crate) const NONCE_LEN: usize = 24;
pub(crate) const COUNTER_LEN: usize = 4;
const _: () = assert!(NONCE_PREFIX_LEN + COUNTER_LEN + 1 == NONCE_LEN);

pub(crate) 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<&SecretBytes32>,
    passphrase: Option<&SecretVec>,
) -> 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()))?;
            raw.with_array(|raw| out.with_mut_array(|out| out.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);
            pw.with_slice(|pw| out.with_mut_array(|out| argon.hash_password_into(pw, salt, out)))?;
        }
    }
    Ok(out)
}

/// Build the AEAD cipher from the protected key. The cipher holds an
/// unprotected copy of the key while alive; `chacha20poly1305` zeroizes that
/// copy on drop. Wrapping in `Arc` lets us share it across worker threads.
fn build_aead(key: &SecretBytes32) -> Arc<XChaCha20Poly1305> {
    Arc::new(key.with_array(|key| XChaCha20Poly1305::new(key.into())))
}

/// Bump the per-chunk counter; surface a domain error on overflow rather than
/// panicking on debug or wrapping in release.
pub(crate) fn bump_counter(counter: u32) -> Result<u32, FcryError> {
    counter
        .checked_add(1)
        .ok_or_else(|| FcryError::Format("STREAM counter overflow (input too large)".into()))
}

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

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

    let flags = if plaintext_length.is_some() {
        FLAG_LENGTH_COMMITTED
    } else {
        0
    };
    let header = Header {
        alg: AlgId::XChaCha20Poly1305,
        flags,
        chunk_size,
        kdf,
        nonce_prefix,
        plaintext_length,
    };
    let aad = Arc::new(header.encode());
    f_encrypted.write_all(&aad)?;

    let aead = build_aead(key);

    if threads > 1 {
        return pipeline::encrypt_parallel(
            f_plain,
            f_encrypted,
            aead,
            aad,
            nonce_prefix,
            chunk_sz,
            threads,
            plaintext_length,
        );
    }

    let mut buf = vec![0u8; chunk_sz];
    let mut counter: u32 = 0;
    let mut bytes_seen: u64 = 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);
                bytes_seen = bytes_seen.saturating_add(chunk_sz as u64);
                counter = bump_counter(counter)?;
            }
            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)?;
                bytes_seen = bytes_seen.saturating_add(n as u64);
                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;
            }
        }
    }

    if let Some(committed) = plaintext_length
        && committed != bytes_seen
    {
        // Defense in depth: the input changed between stat and EOF. The
        // committed length is part of the AEAD AAD, so any decrypter would
        // also surface this, but we prefer to fail before publishing the file.
        return Err(FcryError::Format(format!(
            "input length changed during encryption: committed {committed}, read {bytes_seen}"
        )));
    }

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

pub fn decrypt<S: AsRef<str>>(
    input_file: Option<S>,
    output_file: Option<S>,
    raw_key: Option<&SecretBytes32>,
    passphrase: Option<&SecretVec>,
    threads: usize,
) -> Result<(), FcryError> {
    let mut reader = open_input(input_file)?.reader;
    let header = Header::read(&mut reader)?;
    let aad = Arc::new(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 = build_aead(&key);

    if threads > 1 {
        return pipeline::decrypt_parallel(
            f_encrypted,
            f_plain,
            aead,
            aad,
            header.nonce_prefix,
            cipher_chunk,
            threads,
            header.plaintext_length,
        );
    }

    let mut buf = vec![0u8; cipher_chunk];
    let mut counter: u32 = 0;
    let mut bytes_written: u64 = 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)?;
                bytes_written = bytes_written.saturating_add(buf.len() as u64);
                buf.resize(cipher_chunk, 0);
                counter = bump_counter(counter)?;
            }
            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)?;
                bytes_written = bytes_written.saturating_add(buf.len() as u64);
                break;
            }
            ReadInfoChunk::Empty => {
                return Err(FcryError::Format(
                    "truncated ciphertext: missing final chunk".into(),
                ));
            }
        }
    }

    if let Some(committed) = header.plaintext_length
        && committed != bytes_written
    {
        return Err(FcryError::Format(format!(
            "decrypted length {bytes_written} disagrees with committed {committed}"
        )));
    }

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

/// Random-access decrypt of a byte range. Requires a seekable input file
/// whose header has `FLAG_LENGTH_COMMITTED` set, so we know exactly where
/// each ciphertext chunk lives and which chunk is the last (its nonce uses
/// the STREAM last-block flag).
pub fn decrypt_range<S: AsRef<str>>(
    input_file: &str,
    output_file: Option<S>,
    raw_key: Option<&SecretBytes32>,
    passphrase: Option<&SecretVec>,
    offset: u64,
    length: u64,
) -> Result<(), FcryError> {
    let file = File::open(input_file)?;
    let mut reader = BufReader::new(file);
    let header = Header::read(&mut reader)?;
    let aad = header.encode();
    let header_len = aad.len() as u64;

    let total = header.plaintext_length.ok_or_else(|| {
        FcryError::Format(
            "random-access decrypt requires a length-committed header (encrypt from a regular file)".into(),
        )
    })?;

    let end = offset
        .checked_add(length)
        .ok_or_else(|| FcryError::Format("offset + length overflows u64".into()))?;
    if end > total {
        return Err(FcryError::Format(format!(
            "range [{offset}, {end}) exceeds plaintext length {total}"
        )));
    }

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

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

    // Layout invariants:
    //   n_chunks = ceil(total / chunk_sz), but always ≥ 1 (the empty file
    //   still authenticates a single empty "last" chunk).
    //   last_idx = n_chunks - 1
    //   last_pt  = total - last_idx * chunk_sz   (in [0, chunk_sz])
    let (n_chunks, last_pt) = if total == 0 {
        (1u64, 0u64)
    } else {
        let n = total.div_ceil(chunk_sz);
        let last = total - (n - 1) * chunk_sz;
        (n, last)
    };
    let last_idx = n_chunks - 1;

    let mut out = OutSink::open(output_file)?;

    if length == 0 {
        out.commit()?;
        return Ok(());
    }

    let start_chunk = offset / chunk_sz;
    let end_chunk = (end - 1) / chunk_sz;

    // Reusable buffer sized to a full chunk + tag.
    let mut buf = Vec::with_capacity(cipher_chunk as usize);

    let mut file = reader.into_inner();

    for i in start_chunk..=end_chunk {
        let i_u32 =
            u32::try_from(i).map_err(|_| FcryError::Format("chunk index exceeds u32".into()))?;
        let is_last = i == last_idx;
        let cipher_len = if is_last {
            last_pt + TAG_LEN as u64
        } else {
            cipher_chunk
        };
        let cipher_len_usz =
            usize::try_from(cipher_len).map_err(|_| FcryError::Format("chunk too big".into()))?;

        let chunk_offset = header_len + i * cipher_chunk;
        file.seek(SeekFrom::Start(chunk_offset))?;
        buf.clear();
        buf.resize(cipher_len_usz, 0);
        file.read_exact(&mut buf)?;

        let nonce = make_nonce(&header.nonce_prefix, i_u32, is_last);
        aead.decrypt_in_place(&nonce, &aad, &mut buf)?;

        // `buf` is now plaintext for this chunk. Compute the chunk's plaintext
        // window in absolute bytes and intersect with the requested range.
        let chunk_start = i * chunk_sz;
        let chunk_end = chunk_start + buf.len() as u64;
        let lo = offset.max(chunk_start) - chunk_start;
        let hi = end.min(chunk_end) - chunk_start;
        out.write_all(&buf[lo as usize..hi as usize])?;
    }

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