use std::fmt::Debug;

use chacha20poly1305::{
    ChaCha20Poly1305, Nonce,
    aead::{Aead as _, AeadCore, Generate as _, Payload, array::typenum::Unsigned as _},
};
use subtle::ConstantTimeEq as _;
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

use crate::{
    Error,
    key::{KdfParams, Key},
    secret_buffer::SecretBuffer,
    storage::{FormatVersion, Header, Storage},
};

/// Length of the ChaCha20-Poly1305 authentication tag appended to the
/// ciphertext.
const TAG_LEN: usize = <ChaCha20Poly1305 as AeadCore>::TagSize::USIZE;

/// Trait for types that can be stored in a `CipherBox`.
///
/// Types must be serializable, deserializable, and zeroizable for security.
pub trait CipherBoxContent: Zeroize + serde::Serialize + for<'a> serde::Deserialize<'a> {}
impl<T> CipherBoxContent for T where T: Zeroize + serde::Serialize + for<'a> serde::Deserialize<'a> {}

/// A secure container that encrypts its contents using ChaCha20-Poly1305.
///
/// `CipherBox` provides encrypted storage for any serializable type. The contents
/// are automatically zeroized when the box is dropped for security.
///
/// # Examples
///
/// ```
/// # use cipherbox::CipherBox;
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // Create a box with some data
/// let data = String::from("secret message");
/// let cipher_box = CipherBox::new(data);
///
/// // Access the data through Deref
/// assert_eq!(cipher_box.as_str(), "secret message");
///
/// // Write to encrypted storage
/// let mut buffer = Vec::new();
/// cipher_box.write("password123", &mut buffer)?;
///
/// // Read back from encrypted storage
/// let restored: CipherBox<String> = CipherBox::read("password123", buffer.as_slice())?;
/// assert_eq!(restored.as_str(), "secret message");
/// # Ok(())
/// # }
/// ```
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct CipherBox<T: CipherBoxContent> {
    value: T,
}

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

impl<T> CipherBox<T>
where
    T: CipherBoxContent,
{
    /// Creates a new `CipherBox` containing the given value.
    ///
    /// The value is stored in memory unencrypted until written to storage.
    ///
    /// # Examples
    ///
    /// ```
    /// # use cipherbox::CipherBox;
    /// let secret_data = vec![1, 2, 3, 4, 5];
    /// let cipher_box = CipherBox::new(secret_data);
    /// assert_eq!(cipher_box.len(), 5);
    /// ```
    pub fn new(value: T) -> Self {
        CipherBox { value }
    }

    /// Reads and decrypts a `CipherBox` from a source.
    ///
    /// Attempts to decrypt the data using the provided password. The source must contain
    /// data that was previously written using `CipherBox::write` with the same password.
    ///
    /// # Arguments
    ///
    /// * `password` - The password used to encrypt the data
    /// * `source` - A reader containing the encrypted data
    ///
    /// # Returns
    ///
    /// The decrypted `CipherBox` on success, or an error if decryption fails.
    ///
    /// # Examples
    ///
    /// ```
    /// # use cipherbox::CipherBox;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// // First, create and encrypt some data
    /// let original = CipherBox::new(String::from("hello"));
    /// let mut buffer = Vec::new();
    /// original.write("secret", &mut buffer)?;
    ///
    /// // Then read it back
    /// let restored: CipherBox<String> = CipherBox::read("secret", buffer.as_slice())?;
    /// assert_eq!(restored.as_str(), "hello");
    /// # Ok(())
    /// # }
    /// ```
    pub fn read(
        password: impl AsRef<str>,
        source: impl std::io::Read,
    ) -> Result<CipherBox<T>, Error> {
        let storage = Storage::load(source)?;
        let header = storage.header();
        let key = Key::from_password(password, &header.salt, &header.kdf)?;

        // Verify the key commitment before attempting decryption. This rejects a
        // wrong key up front and gives the construction the key-commitment
        // property that the Poly1305 tag alone cannot provide.
        if !bool::from(key.commitment().as_slice().ct_eq(&header.commitment)) {
            return Err(Error::CommitmentMismatch);
        }

        let plaintext = Zeroizing::new(
            key.cipher()
                .decrypt(
                    &Nonce::from(header.nonce),
                    Payload {
                        msg: storage.data(),
                        aad: &header.encode(),
                    },
                )
                .map_err(|_| Error::Authentication)?,
        );

        drop(storage);

        // Strings and byte strings are borrowed from the plaintext buffer and
        // copied once, at their exact size, into the deserialized value.
        let value =
            cbor4ii::serde::from_slice(&plaintext).map_err(|e| Error::Deserialization(e.into()))?;

        Ok(CipherBox { value })
    }

    /// Encrypts and writes the `CipherBox` to a destination.
    ///
    /// Encrypts the contents with the given password and writes the encrypted
    /// data to the destination. The box is left intact, so its contents can
    /// still be used and written again; drop it to wipe the contents from
    /// memory. Intermediate copies of the contents made while encrypting are
    /// zeroized before this returns.
    ///
    /// # Arguments
    ///
    /// * `password` - The password to use for encryption
    /// * `destination` - A writer to store the encrypted data
    ///
    /// # Examples
    ///
    /// ```
    /// # use cipherbox::CipherBox;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let data = vec![1, 2, 3, 4, 5];
    /// let cipher_box = CipherBox::new(data);
    ///
    /// let mut buffer = Vec::new();
    /// cipher_box.write("my-password", &mut buffer)?;
    ///
    /// // The buffer now contains encrypted data, and the box is still usable
    /// assert!(!buffer.is_empty());
    /// assert_eq!(cipher_box.len(), 5);
    /// # Ok(())
    /// # }
    /// ```
    pub fn write(
        &self,
        password: impl AsRef<str>,
        destination: impl std::io::Write,
    ) -> Result<(), Error> {
        let mut serializer = cbor4ii::serde::Serializer::new(SecretBuffer::new());
        self.value
            .serialize(&mut serializer)
            .map_err(|e| Error::Serialization(e.into()))?;
        let plaintext = serializer.into_inner();

        let mut salt = [0u8; 16];
        rand::fill(&mut salt);
        let kdf = KdfParams::DEFAULT;
        let key = Key::from_password(password, &salt, &kdf)?;

        let header = Header {
            version: FormatVersion::CURRENT,
            kdf,
            salt,
            nonce: Nonce::generate().into(),
            commitment: key.commitment(),
            length: (plaintext.as_slice().len() + TAG_LEN) as u64,
        };

        let ciphertext = key
            .cipher()
            .encrypt(
                &Nonce::from(header.nonce),
                Payload {
                    msg: plaintext.as_slice(),
                    aad: &header.encode(),
                },
            )
            .map_err(|_| Error::Encryption)?;

        drop(plaintext);

        Storage::new(header, &ciphertext).store(destination)
    }
}

impl<T> std::ops::Deref for CipherBox<T>
where
    T: CipherBoxContent,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<T> std::ops::DerefMut for CipherBox<T>
where
    T: CipherBoxContent,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.value
    }
}

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

    #[test]
    fn save_and_keep_working() {
        let mut vault = CipherBox::new(String::from("first draft"));

        let mut first = Vec::new();
        vault.write("just-a-test", &mut first).unwrap();

        vault.push_str(", revised");

        let mut second = Vec::new();
        vault.write("just-a-test", &mut second).unwrap();

        assert_eq!(vault.as_str(), "first draft, revised");

        let restored: CipherBox<String> = CipherBox::read("just-a-test", first.as_slice()).unwrap();
        assert_eq!(restored.as_str(), "first draft");

        let restored: CipherBox<String> =
            CipherBox::read("just-a-test", second.as_slice()).unwrap();
        assert_eq!(restored.as_str(), "first draft, revised");
    }

    #[test]
    fn roundtrip() {
        let mut file = Vec::new();

        let box1 = CipherBox::new(String::from("Hello World"));

        assert_eq!(box1.as_str(), "Hello World");
        assert_eq!(&file, &[]);

        box1.write("just-a-test", &mut file).unwrap();

        assert_ne!(&file, &[]);
        assert!(String::from_utf8_lossy(&file).find("Hello").is_none());

        let box2: CipherBox<String> = CipherBox::read("just-a-test", file.as_slice()).unwrap();

        assert_eq!(box2.as_str(), "Hello World");
    }

    #[test]
    fn validation() {
        let mut file = Vec::new();

        let box1 = CipherBox::new(String::from("Hello World"));

        box1.write("just-a-test", &mut file).unwrap();

        CipherBox::<String>::read("wrong-password", file.as_slice())
            .expect_err("Decryption succeeded with the wrong password");
    }

    #[test]
    fn tampering_detection() {
        let mut file = Vec::new();

        let box1 = CipherBox::new(String::from("Hello World"));
        box1.write("test-password", &mut file).unwrap();

        assert!(CipherBox::<String>::read("test-password", file.as_slice()).is_ok());

        let length_offset = Header::V1_LEN - 8;
        let mut shortened = file.clone();
        let length =
            u64::from_be_bytes(shortened[length_offset..Header::V1_LEN].try_into().unwrap());
        shortened[length_offset..Header::V1_LEN].copy_from_slice(&(length - 1).to_be_bytes());

        let flip = |offset: usize, mask: u8| {
            let mut tampered = file.clone();
            tampered[offset] ^= mask;
            tampered
        };

        type IsExpected = fn(&Error) -> bool;

        // Offsets follow the layout documented on `Header`; each integer's
        // low byte is its last.
        let version = SIGNATURE.len();
        let memory = version + 1;
        let iterations = memory + 4;
        let parallelism = iterations + 4;
        let salt = parallelism + 4;
        let nonce = salt + 16;
        let commitment = nonce + 12;
        let cases: [(&str, Vec<u8>, IsExpected); 9] = [
            ("version", flip(version, 0x02), |e| {
                matches!(e, Error::UnsupportedVersion(3))
            }),
            ("memory cost", flip(memory + 3, 0x01), |e| {
                matches!(e, Error::CommitmentMismatch)
            }),
            ("iterations", flip(iterations + 3, 0x01), |e| {
                matches!(e, Error::CommitmentMismatch)
            }),
            ("parallelism", flip(parallelism + 3, 0x02), |e| {
                matches!(e, Error::CommitmentMismatch)
            }),
            ("salt", flip(salt, 0x01), |e| {
                matches!(e, Error::CommitmentMismatch)
            }),
            ("nonce", flip(nonce, 0x01), |e| {
                matches!(e, Error::Authentication)
            }),
            ("commitment", flip(commitment, 0x01), |e| {
                matches!(e, Error::CommitmentMismatch)
            }),
            ("length", shortened, |e| matches!(e, Error::Authentication)),
            ("ciphertext", flip(Header::V1_LEN, 0x01), |e| {
                matches!(e, Error::Authentication)
            }),
        ];

        for (field, tampered, expected) in cases {
            match CipherBox::<String>::read("test-password", tampered.as_slice()) {
                Err(e) if expected(&e) => {}
                Err(e) => panic!("tampering with the {field} gave an unexpected error: {e:?}"),
                Ok(_) => panic!("tampering with the {field} was not detected"),
            }
        }
    }

    #[test]
    fn wrong_password_fails_commitment_check() {
        let mut file = Vec::new();

        let box1 = CipherBox::new(String::from("Hello World"));
        box1.write("just-a-test", &mut file).unwrap();

        // A wrong password derives a different commitment, so decryption is
        // rejected at the commitment check before the AEAD is ever invoked.
        assert!(matches!(
            CipherBox::<String>::read("wrong-password", file.as_slice()),
            Err(Error::CommitmentMismatch)
        ));
    }
}