use chacha20::{ChaCha20Poly1305, KeyInit, aead::Aead};
use zeroize::{Zeroize, Zeroizing};

use crate::{App, Error, Nonce, Password, Result};

const SALT_LEN: usize = 32;

/// Represents a salt used for password-based key derivation.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Zeroize)]
pub struct Salt([u8; SALT_LEN]);

impl Salt {
    /// Represents the byte length.
    pub const LEN: usize = SALT_LEN;
    pub const ENCRYPTED_LEN: usize = Self::LEN + App::MAC_LEN;

    /// Creates a new [Salt].
    pub const fn new() -> Self {
        Self([0u8; Self::LEN])
    }

    /// Creates a new random [Salt].
    pub fn random() -> Self {
        let mut salt = Salt::new();
        rand::fill(salt.as_mut());
        salt
    }

    /// Converts a slice into a [Salt].
    pub fn from_slice(val: &[u8]) -> Self {
        let mut salt = Salt::new();
        let len = std::cmp::min(Self::LEN, val.len());
        salt.0[..len].copy_from_slice(&val[..len]);
        salt
    }

    /// Gets a reference to the inner slice.
    pub const fn as_slice(&self) -> &[u8] {
        &self.0
    }

    /// Gets the [Salt] length.
    pub const fn len(&self) -> usize {
        Self::LEN
    }

    /// Gets whether the [Salt] is all zeroes.
    pub fn is_empty(&self) -> bool {
        self.0 == [0u8; Self::LEN]
    }

    /// Encrypts the [Salt] with the provided key and nonce.
    pub fn encrypt(
        &self,
        key: &Password,
        nonce: &Nonce,
    ) -> Result<Zeroizing<[u8; Self::ENCRYPTED_LEN]>> {
        let cipher = ChaCha20Poly1305::new_from_slice(key.as_ref())?;

        cipher
            .encrypt(nonce.as_nonce(), self.as_slice())
            .map_err(Error::from)
            .and_then(|v| {
                v.try_into()
                    .map_err(|_| Error::io_crypto("invalid encrypted Salt output length"))
            })
            .map(Zeroizing::new)
    }

    /// Decrypts the provided ciphertext into a [Salt].
    ///
    /// `cipher_text` should be the encrypted [Salt] with the appended Poly1305 MAC.
    ///
    /// **NOTE**: `cipher_text` must be [ENCRYPTED_LEN](Self::ENCRYPTED_LEN) bytes long.
    pub fn decrypt(key: &Password, nonce: &Nonce, cipher_text: &[u8]) -> Result<Self> {
        let len = cipher_text.len();
        if len != Self::ENCRYPTED_LEN {
            Err(Error::io_crypto(format!(
                "invalid encrypted Salt length: {len}"
            )))
        } else {
            let cipher = ChaCha20Poly1305::new_from_slice(key.as_ref())?;

            cipher
                .decrypt(nonce.as_nonce(), cipher_text)
                .map(Zeroizing::new)
                .map(|d| Self::from_slice(d.as_ref()))
                .map_err(Error::from)
        }
    }
}

impl Default for Salt {
    fn default() -> Self {
        Self::new()
    }
}

impl AsRef<[u8]> for Salt {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl AsMut<[u8]> for Salt {
    fn as_mut(&mut self) -> &mut [u8] {
        self.0.as_mut()
    }
}

impl From<&[u8]> for Salt {
    fn from(val: &[u8]) -> Self {
        Self::from_slice(val)
    }
}

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

    const RAND_SALTS: usize = 1024;

    #[test]
    fn test_salt() {
        let mut salts = Vec::with_capacity(RAND_SALTS);
        let key = Password::hash(b"test_salt_key".as_ref(), &Salt::new());

        (0..RAND_SALTS).for_each(|_| {
            let salt = Salt::random();
            let nonce = Nonce::random();

            assert!(salt != Salt::new());
            assert!(!salt.is_empty());

            assert_eq!(Salt::from_slice(salt.as_slice()), salt);

            let cipher_text = salt.encrypt(&key, &nonce).unwrap();
            let decrypted_salt = Salt::decrypt(&key, &nonce, cipher_text.as_ref()).unwrap();

            assert_eq!(decrypted_salt, salt);

            // ensure random salts create unique entries
            assert!(!salts.contains(&salt));

            salts.push(salt);
        });
    }
}
