// Copyright (c) 2026 Christoph Gaffga
// SPDX-License-Identifier: GPL-3.0-only
// https://github.com/cgaffga/phasmcore

//! Cryptographic primitives for payload encryption.
//!
//! Implements a two-tier key derivation scheme using Argon2id:
//!
//! - **Tier 1 (structural)**: Deterministic key derived from passphrase + fixed
//!   salt. Produces `perm_seed` (coefficient permutation) and `hhat_seed`
//!   (STC matrix generation). Both encoder and decoder derive identical keys.
//!
//! - **Tier 2 (encryption)**: AES-256-GCM-SIV key derived from passphrase +
//!   random salt. The random salt is embedded in the payload frame, so the
//!   decoder recovers it from the extracted data.
//!
//! AES-256-GCM-SIV is chosen over AES-256-GCM for its nonce-misuse resistance,
//! which provides an extra safety margin since the nonce is randomly generated
//! and embedded alongside the ciphertext.

use aes_gcm_siv::{Aes256GcmSiv, KeyInit, Nonce};
use aes_gcm_siv::aead::Aead;
use argon2::Argon2;
use zeroize::Zeroizing;
use crate::stego::error::StegoError;

/// Fixed salt for Ghost Tier-1 (structural) key derivation.
/// This is intentionally fixed so the decoder can reproduce perm_seed/hhat_seed
/// from the passphrase alone, before extracting the payload.
const STRUCTURAL_SALT: &[u8; 16] = b"phasm-ghost-v1\0\0";

/// Fixed salt for Armor Tier-1 (structural) key derivation.
/// Different from Ghost so the same passphrase produces different permutations.
const ARMOR_STRUCTURAL_SALT: &[u8; 16] = b"phasm-armor-v1\0\0";

/// Fixed salt for DFT template key derivation (Phase 3 geometry resilience).
/// Independent from structural/armor keys so the template peaks are uncorrelated.
const TEMPLATE_SALT: &[u8; 16] = b"phasm-tmpl-v1\0\0\0";

/// Fixed salt for Fortress structural key derivation (BA-QIM block permutation).
/// Independent from Armor STDM keys so the block order is uncorrelated.
const FORTRESS_STRUCTURAL_SALT: &[u8; 16] = b"phasm-fort-v1\0\0\0";

/// AES-GCM-SIV nonce length in bytes.
pub const NONCE_LEN: usize = 12;
/// Argon2 salt length in bytes.
pub const SALT_LEN: usize = 16;

/// Fixed salt for Shadow layer structural key derivation (repetition coding).
/// Independent from all other keys so shadow permutations are uncorrelated.
const SHADOW_STRUCTURAL_SALT: &[u8; 16] = b"phasm-shdw-v1\0\0\0";

/// Fixed salt for H.264 Phase 3c MVD-domain structural key. Independent from
/// the Ghost structural salt so the MVD-domain permutation + STC matrix do
/// not correlate with the coefficient-domain ones, keeping the two
/// cross-domain flip sets statistically uncorrelated.
const H264_MVD_STRUCTURAL_SALT: &[u8; 16] = b"phasm-h264mvd-v1";

/// Fixed salt for Fortress empty-passphrase optimization.
/// When passphrase is empty, we use this deterministic salt so it doesn't need
/// to be embedded in the frame (saving 16 bytes). The message is still
/// AES-encrypted so the payload looks random for steganalysis resistance.
/// NOT secret — just a constant to feed into AES key derivation.
pub const FORTRESS_EMPTY_SALT: [u8; SALT_LEN] = *b"phasm-fe-salt00\0";

/// Fixed nonce for Fortress empty-passphrase optimization.
/// When passphrase is empty, we use this deterministic nonce so it doesn't
/// need to be embedded in the frame (saving 12 bytes).
/// NOT secret — just a constant to feed into AES-GCM-SIV.
pub const FORTRESS_EMPTY_NONCE: [u8; NONCE_LEN] = *b"ph-fe-nonce\0";

/// Derive the structural key (Tier 1) from a passphrase.
///
/// Returns a 64-byte buffer: first 32 bytes = perm_seed, last 32 bytes = hhat_seed.
/// This key is deterministic given the passphrase so both encoder and decoder agree.
pub fn derive_structural_key(passphrase: &str) -> Result<Zeroizing<[u8; 64]>, StegoError> {
    let mut output = Zeroizing::new([0u8; 64]);
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), STRUCTURAL_SALT, &mut *output)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(output)
}

/// Derive the H.264 Phase 3c MVD-domain structural key.
///
/// Independent from `derive_structural_key` so Phase 3c's two cross-domain
/// STC runs use uncorrelated permutations and HHat matrices. Same 64-byte
/// layout as the main structural key: first 32 = perm_seed, last 32 =
/// hhat_seed.
pub fn derive_h264_mvd_structural_key(
    passphrase: &str,
) -> Result<Zeroizing<[u8; 64]>, StegoError> {
    let mut output = Zeroizing::new([0u8; 64]);
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), H264_MVD_STRUCTURAL_SALT, &mut *output)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(output)
}

/// H.264 Phase I.0.5: derive a per-GOP seed by mixing a master 32-byte seed
/// with `gop_idx` and a domain label using SHA-256.
///
/// The master seed is already passphrase-derived via `derive_structural_key`
/// or `derive_h264_mvd_structural_key` (each Argon2-expensive but only run
/// once per encode/decode), so this per-GOP derivation can be a fast SHA-256
/// pass — the key material is already secret. Deterministic: same master +
/// `gop_idx` + label → same output across iOS / Android / x86_64 / WASM.
///
/// `label` should be a short distinguishing tag (e.g. `b"coeff-perm"`,
/// `b"coeff-hhat"`) so the four per-GOP seeds (perm + hhat × coeff + MVD)
/// are mutually uncorrelated even when they share a master.
pub fn derive_per_gop_seed_from_master(
    master_seed: &[u8; 32],
    gop_idx: u32,
    label: &[u8],
) -> [u8; 32] {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(b"phasm-h264-gop-v1");
    hasher.update(master_seed);
    hasher.update(label);
    hasher.update(gop_idx.to_le_bytes());
    let digest = hasher.finalize();
    let mut out = [0u8; 32];
    out.copy_from_slice(&digest);
    out
}

/// Derive the Armor structural key (Tier 1) from a passphrase.
///
/// Same structure as Ghost but uses a different salt so the same passphrase
/// produces different permutation/spreading seeds.
pub fn derive_armor_structural_key(passphrase: &str) -> Result<Zeroizing<[u8; 64]>, StegoError> {
    let mut output = Zeroizing::new([0u8; 64]);
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), ARMOR_STRUCTURAL_SALT, &mut *output)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(output)
}

/// Derive the DFT template key from a passphrase.
///
/// Returns a 32-byte key used as a ChaCha20 seed for generating template
/// peak positions. Independent from Ghost/Armor structural keys.
pub fn derive_template_key(passphrase: &str) -> Result<[u8; 32], StegoError> {
    let mut output = [0u8; 32];
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), TEMPLATE_SALT, &mut output)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(output)
}

/// Derive the Fortress structural key from a passphrase.
///
/// Returns a 32-byte key used as a ChaCha20 seed for generating block
/// permutation. Independent from Ghost/Armor/Template structural keys.
pub fn derive_fortress_structural_key(passphrase: &str) -> Result<[u8; 32], StegoError> {
    let mut output = [0u8; 32];
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), FORTRESS_STRUCTURAL_SALT, &mut output)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(output)
}

/// Derive the Shadow structural key from a passphrase.
///
/// Returns a 32-byte key used as a ChaCha20 seed for generating position
/// permutation for shadow layer repetition coding. Independent from all
/// other structural keys. Wrapped in `Zeroizing` to prevent key material
/// from lingering in memory after use.
pub fn derive_shadow_structural_key(passphrase: &str) -> Result<Zeroizing<[u8; 32]>, StegoError> {
    let mut output = Zeroizing::new([0u8; 32]);
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), SHADOW_STRUCTURAL_SALT, &mut *output)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(output)
}

/// Derive the AES-256 encryption key (Tier 2) from passphrase + random salt.
pub fn derive_encryption_key(passphrase: &str, salt: &[u8]) -> Result<Zeroizing<[u8; 32]>, StegoError> {
    let mut key = Zeroizing::new([0u8; 32]);
    Argon2::default()
        .hash_password_into(passphrase.as_bytes(), salt, &mut *key)
        .map_err(|_| StegoError::KeyDerivationFailed)?;
    Ok(key)
}

/// Encrypt plaintext with AES-256-GCM-SIV.
///
/// Returns (ciphertext_with_tag, nonce, salt).
/// The ciphertext includes the 16-byte authentication tag appended by AES-GCM-SIV.
pub fn encrypt(plaintext: &[u8], passphrase: &str) -> Result<(Vec<u8>, [u8; NONCE_LEN], [u8; SALT_LEN]), StegoError> {
    use rand::RngCore;
    let mut rng = rand::thread_rng();

    let mut salt = [0u8; SALT_LEN];
    rng.fill_bytes(&mut salt);

    let mut nonce_bytes = [0u8; NONCE_LEN];
    rng.fill_bytes(&mut nonce_bytes);

    let key = derive_encryption_key(passphrase, &salt)?;
    let cipher = Aes256GcmSiv::new_from_slice(&*key).expect("valid key length");
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher.encrypt(nonce, plaintext).expect("AES-GCM-SIV encrypt should not fail");

    Ok((ciphertext, nonce_bytes, salt))
}

/// Encrypt plaintext with AES-256-GCM-SIV using caller-provided salt and nonce.
///
/// Used by the Fortress compact frame path where salt and nonce are fixed
/// constants rather than random values.
pub fn encrypt_with(
    plaintext: &[u8],
    passphrase: &str,
    salt: &[u8; SALT_LEN],
    nonce_bytes: &[u8; NONCE_LEN],
) -> Result<Vec<u8>, StegoError> {
    let key = derive_encryption_key(passphrase, salt)?;
    let cipher = Aes256GcmSiv::new_from_slice(&*key).expect("valid key length");
    let nonce = Nonce::from_slice(nonce_bytes);

    Ok(cipher.encrypt(nonce, plaintext).expect("AES-GCM-SIV encrypt should not fail"))
}

/// Decrypt ciphertext with AES-256-GCM-SIV.
///
/// Returns the plaintext or `StegoError::DecryptionFailed` if the passphrase is wrong
/// or data is corrupted.
pub fn decrypt(
    ciphertext: &[u8],
    passphrase: &str,
    salt: &[u8],
    nonce_bytes: &[u8; NONCE_LEN],
) -> Result<Vec<u8>, StegoError> {
    let key = derive_encryption_key(passphrase, salt)?;
    let cipher = Aes256GcmSiv::new_from_slice(&*key).expect("valid key length");
    let nonce = Nonce::from_slice(nonce_bytes);

    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| StegoError::DecryptionFailed)
}

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

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let msg = b"Hello, steganography!";
        let passphrase = "secret123";

        let (ct, nonce, salt) = encrypt(msg, passphrase).unwrap();
        let pt = decrypt(&ct, passphrase, &salt, &nonce).unwrap();
        assert_eq!(pt, msg);
    }

    #[test]
    fn wrong_passphrase_fails() {
        let msg = b"secret message";
        let (ct, nonce, salt) = encrypt(msg, "correct").unwrap();
        let result = decrypt(&ct, "wrong", &salt, &nonce);
        assert!(matches!(result, Err(StegoError::DecryptionFailed)));
    }

    #[test]
    fn empty_message_works() {
        let msg = b"";
        let passphrase = "pass";
        let (ct, nonce, salt) = encrypt(msg, passphrase).unwrap();
        let pt = decrypt(&ct, passphrase, &salt, &nonce).unwrap();
        assert_eq!(pt, msg.to_vec());
    }

    #[test]
    fn structural_key_deterministic() {
        let a = derive_structural_key("mypass").unwrap();
        let b = derive_structural_key("mypass").unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn structural_key_differs_by_passphrase() {
        let a = derive_structural_key("pass1").unwrap();
        let b = derive_structural_key("pass2").unwrap();
        assert_ne!(a, b);
    }

    #[test]
    fn ghost_and_armor_structural_keys_differ() {
        let ghost = derive_structural_key("same_pass").unwrap();
        let armor = derive_armor_structural_key("same_pass").unwrap();
        assert_ne!(ghost, armor, "Ghost and Armor keys must differ for the same passphrase");
    }

    #[test]
    fn armor_structural_key_deterministic() {
        let a = derive_armor_structural_key("mypass").unwrap();
        let b = derive_armor_structural_key("mypass").unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn template_key_deterministic() {
        let a = derive_template_key("mypass").unwrap();
        let b = derive_template_key("mypass").unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn fortress_key_deterministic() {
        let a = derive_fortress_structural_key("mypass").unwrap();
        let b = derive_fortress_structural_key("mypass").unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn shadow_key_deterministic() {
        let a = derive_shadow_structural_key("mypass").unwrap();
        let b = derive_shadow_structural_key("mypass").unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn shadow_key_differs_from_others() {
        let ghost = derive_structural_key("same_pass").unwrap();
        let armor = derive_armor_structural_key("same_pass").unwrap();
        let fortress = derive_fortress_structural_key("same_pass").unwrap();
        let template = derive_template_key("same_pass").unwrap();
        let shadow = derive_shadow_structural_key("same_pass").unwrap();
        assert_ne!(&ghost[..32], &shadow[..]);
        assert_ne!(&armor[..32], &shadow[..]);
        assert_ne!(&fortress[..], &shadow[..]);
        assert_ne!(&template[..], &shadow[..]);
    }

    #[test]
    fn fortress_key_differs_from_others() {
        let ghost = derive_structural_key("same_pass").unwrap();
        let armor = derive_armor_structural_key("same_pass").unwrap();
        let fortress = derive_fortress_structural_key("same_pass").unwrap();
        let template = derive_template_key("same_pass").unwrap();
        assert_ne!(&ghost[..32], &fortress[..]);
        assert_ne!(&armor[..32], &fortress[..]);
        assert_ne!(&template[..], &fortress[..]);
    }

    #[test]
    fn template_key_differs_from_structural() {
        let ghost = derive_structural_key("same_pass").unwrap();
        let armor = derive_armor_structural_key("same_pass").unwrap();
        let template = derive_template_key("same_pass").unwrap();
        assert_ne!(&ghost[..32], &template[..]);
        assert_ne!(&armor[..32], &template[..]);
    }

    #[test]
    fn encryption_key_differs_by_salt() {
        let key1 = derive_encryption_key("pass", &[0u8; 16]).unwrap();
        let key2 = derive_encryption_key("pass", &[1u8; 16]).unwrap();
        assert_ne!(key1, key2);
    }

    #[test]
    fn ciphertext_differs_per_encryption() {
        // Even with the same plaintext and passphrase, each encryption
        // should produce different ciphertext (due to random salt + nonce).
        let msg = b"same message";
        let (ct1, _, _) = encrypt(msg, "pass").unwrap();
        let (ct2, _, _) = encrypt(msg, "pass").unwrap();
        assert_ne!(ct1, ct2, "repeated encryptions should produce different ciphertext");
    }
}
