use argon2::{Algorithm, Argon2, Block, Params, Version};
use chacha20poly1305::{ChaCha20Poly1305, KeyInit as _};
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

use crate::{error::Error, normalize};

/// BLAKE3 domain-separation context for deriving the ChaCha20-Poly1305
/// encryption subkey from a key's master secret.
const ENCRYPTION_KEY_CONTEXT: &str = "cipherbox 2026-07-07 chacha20-poly1305 encryption subkey";

/// BLAKE3 domain-separation context for deriving a key's commitment value from
/// its master secret.
const COMMITMENT_CONTEXT: &str = "cipherbox 2026-07-07 key commitment";

/// Argon2id cost parameters used to derive a master secret from a password.
///
/// The parameters are stored alongside every ciphertext, so the values used
/// for new data can be changed without affecting the readability of existing
/// data. Parameters read from storage are bounded by the `MAX_*` limits so
/// that a crafted header cannot demand unbounded memory or time.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct KdfParams {
    memory_kib: u32,
    iterations: u32,
    parallelism: u32,
}

impl KdfParams {
    /// The parameters used when writing new data: 19 MiB, 2 iterations,
    /// 1 lane, matching the OWASP password storage recommendation for Argon2id.
    pub(crate) const DEFAULT: KdfParams = KdfParams {
        memory_kib: 19 * 1024,
        iterations: 2,
        parallelism: 1,
    };

    /// The largest memory cost accepted when reading: 4 GiB.
    pub(crate) const MAX_MEMORY_KIB: u32 = 4 * 1024 * 1024;

    /// The largest iteration count accepted when reading.
    pub(crate) const MAX_ITERATIONS: u32 = 64;

    /// The largest degree of parallelism accepted when reading.
    pub(crate) const MAX_PARALLELISM: u32 = 64;

    /// Creates a set of parameters, rejecting any that Argon2 considers
    /// invalid or that exceed the `MAX_*` limits.
    pub(crate) fn new(memory_kib: u32, iterations: u32, parallelism: u32) -> Result<Self, Error> {
        let params = KdfParams {
            memory_kib,
            iterations,
            parallelism,
        };

        if memory_kib > Self::MAX_MEMORY_KIB
            || iterations > Self::MAX_ITERATIONS
            || parallelism > Self::MAX_PARALLELISM
        {
            return Err(Error::UnsupportedKdfParameters);
        }

        params
            .argon2_params()
            .map_err(|_| Error::UnsupportedKdfParameters)?;

        Ok(params)
    }

    /// Argon2 memory cost in KiB.
    pub(crate) fn memory_kib(&self) -> u32 {
        self.memory_kib
    }

    /// Argon2 iteration count.
    pub(crate) fn iterations(&self) -> u32 {
        self.iterations
    }

    /// Argon2 degree of parallelism.
    pub(crate) fn parallelism(&self) -> u32 {
        self.parallelism
    }

    fn argon2_params(&self) -> Result<Params, argon2::Error> {
        Params::new(
            self.memory_kib,
            self.iterations,
            self.parallelism,
            Some(Key::LEN),
        )
    }
}

/// A 32-byte master secret from which the encryption subkey and the key
/// commitment are derived.
///
/// Keys are automatically zeroized when dropped for security.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct Key {
    master: [u8; Key::LEN],
}

impl Key {
    /// Length of the master secret in bytes.
    pub(crate) const LEN: usize = 32;

    /// Creates a key directly from a 32-byte master secret.
    #[allow(dead_code)]
    pub fn from_bytes(master: [u8; Key::LEN]) -> Key {
        Key { master }
    }

    /// Derives a key from a password and salt using Argon2id (version 0x13).
    ///
    /// The password is first normalized to Unicode Normalization Form C, as
    /// recommended for passwords by RFC 8265, so that visually identical
    /// passwords entered on systems that compose characters differently derive
    /// the same key. Argon2's working memory is wiped before it is freed.
    pub fn from_password(
        password: impl AsRef<str>,
        salt: &[u8; 16],
        params: &KdfParams,
    ) -> Result<Key, Error> {
        let password = normalize::nfc(password.as_ref());
        let argon2_params = params
            .argon2_params()
            .map_err(|_| Error::UnsupportedKdfParameters)?;

        let mut memory = Zeroizing::new(Vec::new());
        memory
            .try_reserve_exact(argon2_params.block_count())
            .map_err(|e| Error::KeyDerivation(e.into()))?;
        memory.resize(argon2_params.block_count(), Block::new());

        let mut key = Key {
            master: [0u8; Key::LEN],
        };

        Argon2::new(Algorithm::Argon2id, Version::V0x13, argon2_params)
            .hash_password_into_with_memory(
                password.as_bytes(),
                salt,
                &mut key.master,
                memory.as_mut_slice(),
            )
            .map_err(|e| Error::KeyDerivation(e.into()))?;

        Ok(key)
    }

    /// Creates the ChaCha20-Poly1305 cipher keyed with this key's encryption
    /// subkey.
    ///
    /// The subkey is derived from the master secret through a BLAKE3
    /// domain-separation context distinct from the commitment's, so publishing
    /// the commitment reveals nothing about it.
    pub fn cipher(&self) -> ChaCha20Poly1305 {
        let subkey = Zeroizing::new(blake3::derive_key(ENCRYPTION_KEY_CONTEXT, &self.master));
        ChaCha20Poly1305::new(&(*subkey).into())
    }

    /// Derives the 32-byte key commitment published alongside the ciphertext.
    ///
    /// On decryption a mismatching commitment proves the supplied password is
    /// wrong (or the data was tampered with), giving the construction the
    /// key-commitment property that the Poly1305 tag alone cannot provide:
    /// finding a second key that produces the same commitment would require a
    /// BLAKE3 collision. The value is not secret and is stored in the clear.
    pub fn commitment(&self) -> [u8; 32] {
        blake3::derive_key(COMMITMENT_CONTEXT, &self.master)
    }
}

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

    #[test]
    fn default_params_are_accepted() {
        let params = KdfParams::DEFAULT;
        assert_eq!(
            KdfParams::new(params.memory_kib, params.iterations, params.parallelism).unwrap(),
            params
        );
    }

    #[test]
    fn params_beyond_limits_are_rejected() {
        let d = KdfParams::DEFAULT;
        for (m, t, p) in [
            (KdfParams::MAX_MEMORY_KIB + 1, d.iterations, d.parallelism),
            (d.memory_kib, KdfParams::MAX_ITERATIONS + 1, d.parallelism),
            (d.memory_kib, d.iterations, KdfParams::MAX_PARALLELISM + 1),
            (d.memory_kib, 0, d.parallelism),
            (4, d.iterations, d.parallelism),
        ] {
            assert!(matches!(
                KdfParams::new(m, t, p),
                Err(Error::UnsupportedKdfParameters)
            ));
        }
    }

    #[test]
    fn equivalent_passwords_derive_the_same_key() {
        let salt = [7u8; 16];
        let composed = Key::from_password("caf\u{e9}", &salt, &KdfParams::DEFAULT).unwrap();
        let decomposed = Key::from_password("cafe\u{301}", &salt, &KdfParams::DEFAULT).unwrap();
        assert_eq!(composed.commitment(), decomposed.commitment());
    }
}