use chacha20poly1305::{
    AeadCore as _, ChaCha20Poly1305, KeyInit as _,
    aead::{Aead as _, OsRng, Payload, rand_core::RngCore as _},
};
use hmac::{Hmac, Mac};
use sha2::Sha256;
use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::Error;

const NONCE_BYTES: usize = 12;

/// Identifies a chunk by a keyed hash (HMAC-SHA256) of its plaintext.
///
/// Using a keyed hash rather than a plain digest means identical content
/// deduplicates within a repository, while someone holding the repository
/// cannot test whether it contains a known plaintext.
#[derive(Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct ChunkId([u8; 32]);

impl ChunkId {
    pub(crate) fn compute(id_key: &[u8; 32], plaintext: &[u8]) -> ChunkId {
        let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(id_key)
            .expect("HMAC-SHA256 accepts keys of any length");
        mac.update(plaintext);

        ChunkId(mac.finalize().into_bytes().into())
    }

    pub(crate) fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }

    pub(crate) fn from_bytes(bytes: [u8; 32]) -> ChunkId {
        ChunkId(bytes)
    }

    pub fn to_hex(&self) -> String {
        hex::encode(self.0)
    }

    pub fn from_hex(text: impl AsRef<str>) -> Result<ChunkId, Error> {
        let text = text.as_ref();
        let mut bytes = [0u8; 32];

        hex::decode_to_slice(text, &mut bytes)
            .map_err(|_| Error::InvalidObjectName(text.to_string()))?;

        Ok(ChunkId(bytes))
    }
}

impl std::fmt::Debug for ChunkId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ChunkId({})", self.to_hex())
    }
}

impl std::fmt::Display for ChunkId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.to_hex())
    }
}

/// The secrets protecting a repository.
///
/// Nothing here is ever stored, anywhere: the keys are re-derived
/// idempotently from the passphrase and the repository's public
/// [`Header`] whenever the repository is opened. Consequently the
/// passphrase *is* the key material — it cannot be changed without
/// re-encrypting everything, and losing it loses the repository.
#[derive(Zeroize, ZeroizeOnDrop)]
pub(crate) struct MasterKeys {
    /// ChaCha20-Poly1305 key sealing every object except the header.
    encryption: [u8; 32],
    /// HMAC key from which chunk ids are computed.
    id: [u8; 32],
    /// Key for the header's passphrase verifier.
    verify: [u8; 32],
}

impl MasterKeys {
    /// Derives the key set from a passphrase: Argon2id over the
    /// repository's public salt yields a master secret, and independent
    /// subkeys are split from it by domain-separated HMAC.
    pub(crate) fn derive(password: &str, header: &Header) -> Result<MasterKeys, Error> {
        let argon2 = argon2::Argon2::new(
            argon2::Algorithm::Argon2id,
            argon2::Version::V0x13,
            argon2::Params::new(
                header.argon2_memory_kib,
                header.argon2_iterations,
                header.argon2_parallelism,
                Some(32),
            )
            .map_err(|_| Error::InvalidHeader)?,
        );

        let mut master = [0u8; 32];
        argon2
            .hash_password_into(password.as_bytes(), &header.salt, &mut master)
            .map_err(|_| Error::InvalidHeader)?;

        let keys = MasterKeys {
            encryption: keyed_digest(&master, &[b"beeping encryption key".as_slice()]),
            id: keyed_digest(&master, &[b"beeping chunk id key".as_slice()]),
            verify: keyed_digest(&master, &[b"beeping verify key".as_slice()]),
        };

        master.zeroize();

        Ok(keys)
    }

    /// The public check value proving a later derivation used the same
    /// passphrase *and* that the header it was derived against is the
    /// one the repository was created with. Committing every header
    /// field means a storage-level attacker cannot silently swap in
    /// weaker KDF parameters or different chunker parameters (the
    /// latter would quietly destroy deduplication): any tampering makes
    /// the verifier mismatch, indistinguishable from a wrong passphrase
    /// and equally fatal to the attempt.
    ///
    /// Knowing the verifier helps an attacker no more than any
    /// ciphertext object does: either way they must brute-force the
    /// passphrase through Argon2.
    pub(crate) fn verifier(&self, header: &Header) -> [u8; 32] {
        keyed_digest(
            &self.verify,
            &[
                b"beeping passphrase verifier".as_slice(),
                &header.version.to_be_bytes(),
                &header.salt,
                &header.argon2_memory_kib.to_be_bytes(),
                &header.argon2_iterations.to_be_bytes(),
                &header.argon2_parallelism.to_be_bytes(),
                &header.config.chunk_min.to_be_bytes(),
                &header.config.chunk_avg.to_be_bytes(),
                &header.config.chunk_max.to_be_bytes(),
            ],
        )
    }

    pub(crate) fn id_key(&self) -> &[u8; 32] {
        &self.id
    }
}

/// A keyed digest over several byte strings, used for deriving
/// deterministic identifiers. Each part's length is mixed in, so distinct
/// part lists never collide by concatenation.
pub(crate) fn keyed_digest(key: &[u8; 32], parts: &[&[u8]]) -> [u8; 32] {
    let mut mac =
        <Hmac<Sha256> as Mac>::new_from_slice(key).expect("HMAC-SHA256 accepts keys of any length");

    for part in parts {
        mac.update(&(part.len() as u64).to_be_bytes());
        mac.update(part);
    }

    mac.finalize().into_bytes().into()
}

/// Parameters that must stay constant for the life of a repository so
/// that identical content keeps producing identical chunks.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(crate) struct RepoConfig {
    pub chunk_min: u32,
    pub chunk_avg: u32,
    pub chunk_max: u32,
}

impl Default for RepoConfig {
    fn default() -> Self {
        RepoConfig {
            chunk_min: 256 * 1024,
            chunk_avg: 1024 * 1024,
            chunk_max: 4 * 1024 * 1024,
        }
    }
}

/// The repository's public, write-once `header` object.
///
/// It deliberately contains no secret material — only what a client
/// needs to re-derive the keys from the passphrase (salt and KDF
/// parameters), detect a wrong passphrase (the verifier), and chunk
/// content compatibly with every other client of the repository.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(crate) struct Header {
    pub version: u16,
    pub salt: [u8; 16],
    pub argon2_memory_kib: u32,
    pub argon2_iterations: u32,
    pub argon2_parallelism: u32,
    pub verifier: [u8; 32],
    pub config: RepoConfig,
}

pub(crate) const HEADER_VERSION: u16 = 1;

impl Header {
    /// A fresh header with a random salt and current cost parameters.
    /// The verifier starts zeroed and is filled in after the first key
    /// derivation.
    pub(crate) fn generate() -> Header {
        let mut salt = [0u8; 16];
        OsRng.fill_bytes(&mut salt);

        let defaults = argon2::Params::DEFAULT;

        Header {
            version: HEADER_VERSION,
            salt,
            argon2_memory_kib: defaults.m_cost(),
            argon2_iterations: defaults.t_cost(),
            argon2_parallelism: defaults.p_cost(),
            verifier: [0u8; 32],
            config: RepoConfig::default(),
        }
    }
}

/// Encrypts `plaintext` with the repository's master key, producing
/// `nonce || ciphertext`. The AAD binds the result to the object it names,
/// so an object substituted for another fails to open.
pub(crate) fn seal(keys: &MasterKeys, aad: &[u8], plaintext: &[u8]) -> Result<Vec<u8>, Error> {
    let cipher = ChaCha20Poly1305::new((&keys.encryption).into());
    let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);

    let ciphertext = cipher
        .encrypt(
            &nonce,
            Payload {
                msg: plaintext,
                aad,
            },
        )
        .map_err(|_| Error::Crypto)?;

    let mut sealed = Vec::with_capacity(NONCE_BYTES + ciphertext.len());
    sealed.extend_from_slice(&nonce);
    sealed.extend_from_slice(&ciphertext);

    Ok(sealed)
}

/// Reverses [`seal`], authenticating the data and its AAD in the process.
pub(crate) fn open(keys: &MasterKeys, aad: &[u8], sealed: &[u8]) -> Result<Vec<u8>, Error> {
    if sealed.len() < NONCE_BYTES {
        return Err(Error::Crypto);
    }

    let (nonce, ciphertext) = sealed.split_at(NONCE_BYTES);
    let cipher = ChaCha20Poly1305::new((&keys.encryption).into());

    cipher
        .decrypt(
            nonce.into(),
            Payload {
                msg: ciphertext,
                aad,
            },
        )
        .map_err(|_| Error::Crypto)
}

/// Defines a 16-byte object identifier: a newtype over the bytes with
/// hex parsing, display, and the derives every one of them needs.
///
/// Snapshots, prune records, and locks are all named this way, and the
/// distinct types keep one from being passed where another belongs.
macro_rules! object_id {
    ($name:ident, $doc:expr) => {
        #[doc = $doc]
        #[derive(Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
        pub struct $name([u8; 16]);

        impl $name {
            pub(crate) fn from_bytes(bytes: [u8; 16]) -> $name {
                $name(bytes)
            }

            pub fn to_hex(&self) -> String {
                hex::encode(self.0)
            }

            pub fn from_hex(text: impl AsRef<str>) -> Result<$name, crate::Error> {
                let text = text.as_ref();
                let mut bytes = [0u8; 16];

                hex::decode_to_slice(text, &mut bytes)
                    .map_err(|_| crate::Error::InvalidObjectName(text.to_string()))?;

                Ok($name(bytes))
            }

            pub(crate) fn as_bytes(&self) -> &[u8; 16] {
                &self.0
            }
        }

        impl std::fmt::Debug for $name {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{}({})", stringify!($name), self.to_hex())
            }
        }

        impl std::fmt::Display for $name {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.write_str(&self.to_hex())
            }
        }
    };
}

/// Gives an id type a source of fresh, unpredictable values.
///
/// Only for the objects named by nothing but themselves. A snapshot's id
/// is derived from what it contains instead, so that a backup redone
/// after an interruption converges on one object rather than piling up
/// duplicates — which is why this is not part of [`object_id`].
macro_rules! random_id {
    ($name:ident) => {
        impl $name {
            pub(crate) fn generate() -> $name {
                let mut bytes = [0u8; 16];
                OsRng.fill_bytes(&mut bytes);

                $name::from_bytes(bytes)
            }
        }
    };
}

object_id!(
    SnapshotId,
    "Identifies a snapshot record within a repository."
);

object_id!(
    PruneId,
    "Identifies the record one prune leaves behind. Their ids, taken \
     together, are what tells a local cache whether anything has been \
     removed from the repository since it was written."
);

object_id!(
    LockId,
    "Identifies one running backup's lock. A lock means the run has \
     references that no snapshot names yet, so pruning would sweep them."
);

random_id!(PruneId);
random_id!(LockId);