//! Argon2id key derivation and XChaCha20-Poly1305 authenticated encryption.
//!
//! This module implements the low-level "encrypt bytes / decrypt bytes" layer.
//! Higher-level typed wrappers (encrypt_entry, encrypt_manifest) live in [`crate::vault`].
//!
//! ## Why XChaCha20-Poly1305 over AES-GCM
//!
//! - **192-bit nonce** (vs. 96-bit for AES-GCM): eliminates nonce collision risk
//!   even with random nonces across billions of encryptions. With AES-GCM's 96-bit
//!   nonce, birthday-bound collisions become probable around 2^48 messages under
//!   the same key -- a real concern for a long-lived vault.
//! - **Fast on WASM and ARM without AES-NI**: ChaCha20 is a pure arithmetic cipher
//!   (add/rotate/XOR) with no dependency on hardware AES acceleration. AES-GCM is
//!   fast *only* with AES-NI; without it, software AES is both slow and vulnerable
//!   to cache-timing side channels.
//!
//! ## Binary ciphertext format
//!
//! Every encrypted blob produced by [`encrypt`] has this layout:
//!
//! ```text
//! [version: 1 byte] [nonce: 24 bytes] [ciphertext + Poly1305 tag: variable]
//! ```
//!
//! - **Version byte** (`0x02`): allows future format changes without ambiguity.
//!   Decryption rejects any version it does not recognize.
//! - **Nonce** (24 bytes): randomly generated per encryption via [`OsRng`].
//!   Stored alongside the ciphertext so the decryptor does not need out-of-band
//!   nonce management.
//! - **Ciphertext + tag**: the AEAD output. The Poly1305 tag (16 bytes) is
//!   appended by the cipher implementation; we do not separate it.
//!
//! ## KDF pipeline
//!
//! [`derive_master_key`] concatenates the passphrase and image_secret as a single
//! password input to Argon2id:
//!
//! ```text
//! password = passphrase_bytes || image_secret (32 bytes)
//! master_key = Argon2id(password, salt, params) -> 32 bytes
//! ```
//!
//! Both factors contribute to the derived key -- compromising one without the
//! other is insufficient. The salt is vault-specific and stored in `.relicario/salt`.

use argon2::{Algorithm, Argon2, Params, Version};
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    XChaCha20Poly1305, XNonce,
};
use rand::{rngs::OsRng, RngCore};
use serde::{Deserialize, Serialize};
use unicode_normalization::UnicodeNormalization;
use zeroize::Zeroizing;

use crate::error::{RelicarioError, Result};

/// Current binary format version. Increment this if the ciphertext layout changes.
pub const VERSION_BYTE: u8 = 0x02;

/// XChaCha20-Poly1305 nonce length: 192 bits = 24 bytes.
const NONCE_LEN: usize = 24;

/// Poly1305 authentication tag length: 128 bits = 16 bytes.
/// Used only for minimum-length validation during decryption.
const TAG_LEN: usize = 16;

/// Total header size: version byte + nonce. The ciphertext (including tag)
/// follows immediately after the header.
const HEADER_LEN: usize = 1 + NONCE_LEN; // version + nonce

/// Encrypt arbitrary plaintext bytes under a 256-bit key using XChaCha20-Poly1305.
///
/// Returns the binary blob in the format: `version(1) || nonce(24) || ciphertext+tag`.
/// A fresh random nonce is generated for each call via the OS CSPRNG.
///
/// # Errors
///
/// Returns [`RelicarioError::Encrypt`] if the underlying AEAD operation fails
/// (extremely unlikely in practice).
pub fn encrypt(key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>> {
    let cipher = XChaCha20Poly1305::new(key.into());

    // Generate a fresh random 24-byte nonce for every encryption.
    // With 192 bits of randomness, nonce reuse probability is negligible
    // even across billions of encryptions under the same key.
    let mut nonce_bytes = [0u8; NONCE_LEN];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = XNonce::from(nonce_bytes);

    let ciphertext = cipher
        .encrypt(&nonce, plaintext)
        .map_err(|e| RelicarioError::Encrypt(e.to_string()))?;

    // Output: version(1) || nonce(24) || ciphertext+tag
    let mut output = Vec::with_capacity(HEADER_LEN + ciphertext.len());
    output.push(VERSION_BYTE);
    output.extend_from_slice(&nonce_bytes);
    output.extend_from_slice(&ciphertext);

    Ok(output)
}

/// Decrypt a blob produced by [`encrypt`], returning the original plaintext.
///
/// Validates the version byte and minimum blob length before attempting
/// authenticated decryption. If the key is wrong or the data has been
/// tampered with, the Poly1305 tag verification fails and [`RelicarioError::Decrypt`]
/// is returned -- with no information about which bytes were wrong (preventing
/// padding oracle / chosen-ciphertext attacks).
///
/// # Errors
///
/// - [`RelicarioError::Format`] if the data is too short or has an unknown version byte.
/// - [`RelicarioError::Decrypt`] if the AEAD tag verification fails (wrong key or
///   tampered data).
pub fn decrypt(key: &[u8; 32], data: &[u8]) -> Result<Vec<u8>> {
    // Minimum valid blob: 1 (version) + 24 (nonce) + 16 (tag) = 41 bytes.
    // A zero-length plaintext produces exactly 41 bytes of output.
    if data.len() < HEADER_LEN + TAG_LEN {
        return Err(RelicarioError::Format(
            "data too short to be valid ciphertext".into(),
        ));
    }

    let found = data[0];
    if found != VERSION_BYTE {
        return Err(RelicarioError::UnsupportedFormatVersion {
            found,
            expected: VERSION_BYTE,
        });
    }

    let nonce = XNonce::from_slice(&data[1..1 + NONCE_LEN]);
    let ciphertext = &data[HEADER_LEN..];

    let cipher = XChaCha20Poly1305::new(key.into());
    let plaintext = cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| RelicarioError::Decrypt)?;

    Ok(plaintext)
}

/// Tunable parameters for the Argon2id key derivation function.
///
/// These are stored in the vault's `.relicario/params.json` so that every client
/// derives the same master key from the same inputs. Making them configurable
/// lets tests use fast params (m=256, t=1, p=1) while production uses strong
/// params (m=64MiB, t=3, p=4).
///
/// The parameters follow Argon2id naming conventions:
/// - `argon2_m`: memory cost in KiB
/// - `argon2_t`: time cost (number of iterations)
/// - `argon2_p`: parallelism degree (number of lanes)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KdfParams {
    /// Memory cost in KiB. Default is 65536 (64 MiB), which makes GPU/ASIC
    /// brute-force attacks expensive. Tests use 256 KiB for speed.
    pub argon2_m: u32,
    /// Time cost (iteration count). Default is 3. Higher values increase CPU
    /// time linearly. Combined with high memory cost, this makes each key
    /// derivation take ~1 second on modern hardware.
    pub argon2_t: u32,
    /// Parallelism degree. Default is 4. Sets the number of independent lanes
    /// in the Argon2id memory-hard computation.
    pub argon2_p: u32,
}

/// Production-strength default parameters: 64 MiB memory, 3 iterations, 4 lanes.
///
/// These are calibrated to take roughly 0.5-1 second on a modern desktop CPU,
/// making brute-force attacks impractical while keeping interactive unlock fast
/// enough for daily use.
impl Default for KdfParams {
    fn default() -> Self {
        Self {
            argon2_m: 65536,
            argon2_t: 3,
            argon2_p: 4,
        }
    }
}

/// Derive a 256-bit master key from the user's passphrase and reference image secret.
///
/// The two factors (passphrase + image_secret) are concatenated into a single
/// password input to Argon2id. This means both factors contribute entropy to
/// the derived key -- compromising one factor alone is insufficient.
///
/// # Arguments
///
/// - `passphrase`: the user's passphrase as raw UTF-8 bytes.
/// - `image_secret`: the 32-byte secret extracted from the reference JPEG via
///   [`crate::imgsecret::extract`].
/// - `salt`: a 32-byte vault-specific salt (stored in `.relicario/salt`).
/// - `params`: the Argon2id tuning parameters (stored in `.relicario/params.json`).
///
/// # Returns
///
/// A 32-byte master key suitable for use with [`encrypt`] and [`decrypt`].
///
/// # Errors
///
/// Returns [`RelicarioError::Kdf`] if the Argon2id parameters are invalid (e.g.,
/// memory cost below the library's minimum).
pub fn derive_master_key(
    passphrase: &[u8],
    image_secret: &[u8; 32],
    salt: &[u8; 32],
    params: &KdfParams,
) -> Result<Zeroizing<[u8; 32]>> {
    let argon2_params = Params::new(
        params.argon2_m,
        params.argon2_t,
        params.argon2_p,
        Some(32),
    )
    .map_err(|e| RelicarioError::Kdf(e.to_string()))?;

    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, argon2_params);

    // Normalize passphrase to NFC. Invalid UTF-8 bytes pass through unchanged.
    let nfc_passphrase: Vec<u8> = match std::str::from_utf8(passphrase) {
        Ok(s) => s.nfc().collect::<String>().into_bytes(),
        Err(_) => passphrase.to_vec(),
    };

    // Length-prefixed concatenation: [u64_be(len(passphrase))][passphrase]
    //                                 [u64_be(32)][image_secret]
    // Eliminates the (passphrase, image_secret) boundary ambiguity (audit H1).
    let mut password = Zeroizing::new(Vec::with_capacity(8 + nfc_passphrase.len() + 8 + 32));
    password.extend_from_slice(&(nfc_passphrase.len() as u64).to_be_bytes());
    password.extend_from_slice(&nfc_passphrase);
    password.extend_from_slice(&32u64.to_be_bytes());
    password.extend_from_slice(image_secret);

    let mut output = Zeroizing::new([0u8; 32]);
    argon2
        .hash_password_into(password.as_slice(), salt, output.as_mut())
        .map_err(|e| RelicarioError::Kdf(e.to_string()))?;

    Ok(output)
}

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

    fn fast_params() -> KdfParams {
        KdfParams {
            argon2_m: 256,
            argon2_t: 1,
            argon2_p: 1,
        }
    }

    #[test]
    fn derive_master_key_deterministic() {
        let passphrase = b"test-passphrase";
        let image_secret = [0x42u8; 32];
        let salt = [0x01u8; 32];
        let params = fast_params();

        let key1 = derive_master_key(passphrase, &image_secret, &salt, &params).unwrap();
        let key2 = derive_master_key(passphrase, &image_secret, &salt, &params).unwrap();

        assert_eq!(*key1, *key2);
    }

    #[test]
    fn derive_master_key_different_passphrase() {
        let image_secret = [0x42u8; 32];
        let salt = [0x01u8; 32];
        let params = fast_params();

        let key1 = derive_master_key(b"passphrase-one", &image_secret, &salt, &params).unwrap();
        let key2 = derive_master_key(b"passphrase-two", &image_secret, &salt, &params).unwrap();

        assert_ne!(*key1, *key2);
    }

    #[test]
    fn derive_master_key_different_image_secret() {
        let passphrase = b"test-passphrase";
        let salt = [0x01u8; 32];
        let params = fast_params();

        let image_secret1 = [0x11u8; 32];
        let image_secret2 = [0x22u8; 32];

        let key1 = derive_master_key(passphrase, &image_secret1, &salt, &params).unwrap();
        let key2 = derive_master_key(passphrase, &image_secret2, &salt, &params).unwrap();

        assert_ne!(*key1, *key2);
    }

    #[test]
    fn encrypt_decrypt_round_trip() {
        let key = [0xABu8; 32];
        let plaintext = b"hello, relicario!";

        let ciphertext = encrypt(&key, plaintext).unwrap();
        let decrypted = decrypt(&key, &ciphertext).unwrap();

        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn decrypt_wrong_key_fails() {
        let key = [0xABu8; 32];
        let wrong_key = [0xCDu8; 32];
        let plaintext = b"sensitive data";

        let ciphertext = encrypt(&key, plaintext).unwrap();
        let result = decrypt(&wrong_key, &ciphertext);

        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), RelicarioError::Decrypt));
    }

    #[test]
    fn decrypt_tampered_data_fails() {
        let key = [0xABu8; 32];
        let plaintext = b"sensitive data";

        let mut ciphertext = encrypt(&key, plaintext).unwrap();
        // Flip a byte in the ciphertext portion (after header)
        let flip_pos = HEADER_LEN + 2;
        ciphertext[flip_pos] ^= 0xFF;

        let result = decrypt(&key, &ciphertext);
        assert!(result.is_err());
    }

    #[test]
    fn ciphertext_format_has_correct_structure() {
        let key = [0x11u8; 32];
        let plaintext = b"test plaintext for structure check";

        let ciphertext = encrypt(&key, plaintext).unwrap();

        // Expected length: 1 (version) + 24 (nonce) + plaintext_len + 16 (tag)
        let expected_len = 1 + 24 + plaintext.len() + 16;
        assert_eq!(ciphertext.len(), expected_len);

        // Version byte must be 0x02
        assert_eq!(ciphertext[0], 0x02);
    }

    #[test]
    fn length_prefix_eliminates_concatenation_ambiguity() {
        // Without length-prefix: ("abc", [0x44, ...]) and ("abcD", [...]) could collide.
        // With length-prefix: distinct inputs always yield distinct keys.
        let salt = [0u8; 32];
        let params = fast_params();

        // Pair A: passphrase "abc", image_secret starts with 0x44
        let mut img_a = [0u8; 32]; img_a[0] = 0x44;
        let key_a = derive_master_key(b"abc", &img_a, &salt, &params).unwrap();

        // Pair B: passphrase "abcD" (one extra char), image_secret starts with original byte 1
        let mut img_b = [0u8; 32]; img_b[0] = 0x44;  // same image
        let key_b = derive_master_key(b"abcD", &img_b, &salt, &params).unwrap();

        // With length-prefix, the keys MUST differ.
        assert_ne!(*key_a, *key_b);
    }

    #[test]
    fn nfc_normalization_collapses_unicode_forms() {
        // "café" can be written as NFC (é = U+00E9) or NFD (e + U+0301).
        // Both must produce the same key after NFC normalization.
        let salt = [0u8; 32];
        let img = [0u8; 32];
        let params = fast_params();

        let nfc = "caf\u{00e9}".as_bytes();  // é precomposed
        let nfd = "cafe\u{0301}".as_bytes(); // e + combining acute

        let key_nfc = derive_master_key(nfc, &img, &salt, &params).unwrap();
        let key_nfd = derive_master_key(nfd, &img, &salt, &params).unwrap();

        assert_eq!(*key_nfc, *key_nfd);
    }

    #[test]
    fn master_key_is_zeroized_on_drop() {
        // Smoke test: master_key returns a Zeroizing<[u8; 32]>, which compiles only if
        // we wrap correctly. The drop wipe is verified by the zeroize crate's tests.
        let salt = [0u8; 32];
        let img = [0u8; 32];
        let params = fast_params();
        let key: zeroize::Zeroizing<[u8; 32]> = derive_master_key(b"x", &img, &salt, &params).unwrap();
        assert_eq!(key.len(), 32);
    }

    #[test]
    fn version_byte_is_0x02() {
        assert_eq!(VERSION_BYTE, 0x02);
    }

    #[test]
    fn decrypt_rejects_v1_blob_with_typed_error() {
        // Construct a v1-style blob: [0x01][24 nonce bytes][16 tag bytes].
        let mut blob = vec![0x01u8];
        blob.extend_from_slice(&[0u8; 24]);
        blob.extend_from_slice(&[0u8; 16]);

        let key = Zeroizing::new([0u8; 32]);
        let err = decrypt(&*key, &blob).expect_err("v1 blob should fail decrypt");
        match err {
            RelicarioError::UnsupportedFormatVersion { found, expected } => {
                assert_eq!(found, 0x01);
                assert_eq!(expected, 0x02);
            }
            other => panic!("expected UnsupportedFormatVersion, got {:?}", other),
        }
    }
}