//! ChaCha20-Poly1305 AEAD cipher suite for DTLS 1.3.
//!
//! This module is compiled out when the `fips140` feature is enabled,
//! since ChaCha20 is not a FIPS 140 approved algorithm.

use dimpl::crypto::{Aad, Buf, Cipher, Dtls13CipherSuite, HashAlgorithm, Nonce};
use dimpl::crypto::{SupportedDtls13CipherSuite, TmpBuf};

use openssl::cipher_ctx::CipherCtx;

use super::cipher_suite::{aead_decrypt, aead_encrypt};

const CHACHA20_POLY1305_TAG_LEN: usize = 16;
const CHACHA20_POLY1305_KEY_LEN: usize = 32;
const CHACHA20_POLY1305_IV_LEN: usize = 12;

/// ChaCha20-Poly1305 cipher implementation using OpenSSL.
struct ChaCha20Poly1305 {
    key: [u8; CHACHA20_POLY1305_KEY_LEN],
}

impl std::fmt::Debug for ChaCha20Poly1305 {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ChaCha20Poly1305").finish_non_exhaustive()
    }
}

impl ChaCha20Poly1305 {
    fn new(key: &[u8]) -> Result<Self, String> {
        let key: [u8; CHACHA20_POLY1305_KEY_LEN] = key
            .try_into()
            .map_err(|_| format!("Invalid key size for ChaCha20-Poly1305: {}", key.len()))?;
        Ok(Self { key })
    }
}

impl Drop for ChaCha20Poly1305 {
    fn drop(&mut self) {
        for b in self.key.iter_mut() {
            // SAFETY: Volatile write prevents the compiler from eliding this zeroing.
            unsafe { std::ptr::write_volatile(b, 0) };
        }
    }
}

impl Cipher for ChaCha20Poly1305 {
    fn encrypt(&mut self, plaintext: &mut Buf, aad: Aad, nonce: Nonce) -> Result<(), String> {
        aead_encrypt(
            openssl::cipher::Cipher::chacha20_poly1305(),
            &self.key,
            plaintext,
            aad,
            nonce,
            CHACHA20_POLY1305_TAG_LEN,
        )
    }

    fn decrypt(&mut self, ciphertext: &mut TmpBuf, aad: Aad, nonce: Nonce) -> Result<(), String> {
        aead_decrypt(
            openssl::cipher::Cipher::chacha20_poly1305(),
            &self.key,
            ciphertext,
            aad,
            nonce,
            CHACHA20_POLY1305_TAG_LEN,
        )
    }
}

/// TLS_CHACHA20_POLY1305_SHA256 cipher suite (TLS 1.3 / DTLS 1.3).
#[derive(Debug)]
pub(super) struct Tls13ChaCha20Poly1305Sha256;

impl SupportedDtls13CipherSuite for Tls13ChaCha20Poly1305Sha256 {
    fn suite(&self) -> Dtls13CipherSuite {
        Dtls13CipherSuite::CHACHA20_POLY1305_SHA256
    }

    fn hash_algorithm(&self) -> HashAlgorithm {
        HashAlgorithm::SHA256
    }

    fn key_len(&self) -> usize {
        CHACHA20_POLY1305_KEY_LEN
    }

    fn iv_len(&self) -> usize {
        CHACHA20_POLY1305_IV_LEN
    }

    fn tag_len(&self) -> usize {
        CHACHA20_POLY1305_TAG_LEN
    }

    fn create_cipher(&self, key: &[u8]) -> Result<Box<dyn Cipher>, String> {
        Ok(Box::new(ChaCha20Poly1305::new(key)?))
    }

    fn encrypt_sn(&self, sn_key: &[u8], sample: &[u8; 16]) -> [u8; 16] {
        if sn_key.len() != 32 {
            panic!(
                "encrypt_sn: invalid ChaCha20 key length {} (expected 32)",
                sn_key.len()
            );
        }
        // RFC 9147 Section 4.2.3 / RFC 9001 Section 5.4.4: For ChaCha20-Poly1305,
        // the mask is generated by treating the sample as a nonce for ChaCha20
        // with a zero block counter and encrypting zero bytes.
        let cipher = openssl::cipher::Cipher::chacha20();
        let mut ctx = CipherCtx::new().expect("CipherCtx::new");
        ctx.encrypt_init(Some(cipher), Some(sn_key), Some(sample))
            .expect("encrypt_init");

        let mut output = [0u8; 32];
        let input = [0u8; 16];
        let count = ctx
            .cipher_update(&input, Some(&mut output))
            .expect("cipher_update");
        let final_count = ctx
            .cipher_final(&mut output[count..])
            .expect("cipher_final");
        debug_assert_eq!(
            count + final_count,
            16,
            "ChaCha20 stream cipher should not pad"
        );

        let mut result = [0u8; 16];
        result.copy_from_slice(&output[..16]);
        result
    }
}

pub(super) static TLS13_CHACHA20_POLY1305_SHA256: Tls13ChaCha20Poly1305Sha256 =
    Tls13ChaCha20Poly1305Sha256;

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

    use crate::dimpl_provider::test_utils::hex_to_vec as hex;

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let key = [0x42u8; 32];
        let nonce = Nonce([0x01u8; 12]);
        let plaintext = b"hello world, this is a test for ChaCha20-Poly1305";

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        // Encrypt
        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);
        cipher
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce)
            .unwrap();

        // Ciphertext should be plaintext_len + 16 (tag)
        assert_eq!(buf.len(), plaintext.len() + CHACHA20_POLY1305_TAG_LEN);
        assert_ne!(&buf.as_ref()[..plaintext.len()], &plaintext[..]);

        // Decrypt
        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        cipher
            .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
            .unwrap();
        assert_eq!(tmp.as_ref(), plaintext);
    }

    #[test]
    fn wrong_key_fails_decrypt() {
        let key1 = [0x42u8; 32];
        let key2 = [0x43u8; 32];
        let nonce = Nonce([0x01u8; 12]);
        let plaintext = b"secret";

        let mut cipher1 = ChaCha20Poly1305::new(&key1).unwrap();
        let mut cipher2 = ChaCha20Poly1305::new(&key2).unwrap();

        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);
        cipher1
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce)
            .unwrap();

        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        assert!(
            cipher2
                .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
                .is_err()
        );
    }

    #[test]
    fn invalid_key_size_rejected() {
        assert!(ChaCha20Poly1305::new(&[0u8; 16]).is_err());
        assert!(ChaCha20Poly1305::new(&[0u8; 31]).is_err());
        assert!(ChaCha20Poly1305::new(&[0u8; 32]).is_ok());
    }

    /// RFC 8439 Section 2.8.2 — AEAD construction test vector.
    #[test]
    fn rfc8439_aead_test_vector() {
        let key = hex("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f");
        let nonce_bytes = hex("070000004041424344454647");
        let aad_bytes = hex("50515253c0c1c2c3c4c5c6c7");
        let plaintext = b"Ladies and Gentlemen of the class of '99: \
If I could offer you only one tip for the future, sunscreen would be it.";

        let expected_ciphertext = hex("d31a8d34648e60db7b86afbc53ef7ec2\
             a4aded51296e08fea9e2b5a736ee62d6\
             3dbea45e8ca9671282fafb69da92728b\
             1a71de0a9e060b2905d6a5b67ecd3b36\
             92ddbd7f2d778b8c9803aee328091b58\
             fab324e4fad675945585808b4831d7bc\
             3ff4def08e4b7a9de576d26586cec64b\
             6116");
        let expected_tag = hex("1ae10b594f09e26a7e902ecbd0600691");

        let nonce = Nonce(nonce_bytes.as_slice().try_into().unwrap());
        let mut aad_arr = arrayvec::ArrayVec::<u8, 13>::new();
        aad_arr.try_extend_from_slice(&aad_bytes).unwrap();

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        // Encrypt
        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);
        cipher.encrypt(&mut buf, Aad(aad_arr), nonce).unwrap();

        // Verify ciphertext (excluding tag)
        let ct_len = buf.len() - CHACHA20_POLY1305_TAG_LEN;
        assert_eq!(&buf.as_ref()[..ct_len], &expected_ciphertext[..]);
        // Verify tag
        assert_eq!(&buf.as_ref()[ct_len..], &expected_tag[..]);

        // Verify decrypt roundtrip
        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        let mut aad_arr2 = arrayvec::ArrayVec::<u8, 13>::new();
        aad_arr2.try_extend_from_slice(&aad_bytes).unwrap();
        cipher.decrypt(&mut tmp, Aad(aad_arr2), nonce).unwrap();
        assert_eq!(tmp.as_ref(), plaintext);
    }

    /// Verify that modifying the AAD causes decryption to fail.
    #[test]
    fn aad_tamper_detected() {
        let key = [0x42u8; 32];
        let nonce = Nonce([0x01u8; 12]);
        let plaintext = b"authenticated data test";

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);

        cipher
            .encrypt(&mut buf, Aad([0x00u8; 13].into()), nonce)
            .unwrap();

        // Tamper with AAD
        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        assert!(
            cipher
                .decrypt(&mut tmp, Aad([0x01u8; 13].into()), nonce)
                .is_err()
        );
    }

    /// Verify that a wrong nonce causes decryption to fail.
    #[test]
    fn wrong_nonce_fails() {
        let key = [0x42u8; 32];
        let nonce1 = Nonce([0x01u8; 12]);
        let nonce2 = Nonce([0x02u8; 12]);
        let plaintext = b"nonce test";

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);
        cipher
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce1)
            .unwrap();

        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        assert!(
            cipher
                .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce2)
                .is_err()
        );
    }

    /// Corrupted ciphertext tag byte should cause decryption failure.
    #[test]
    fn tag_corruption_detected() {
        let key = [0x42u8; 32];
        let nonce = Nonce([0x01u8; 12]);
        let plaintext = b"tag corruption test";

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);
        cipher
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce)
            .unwrap();

        // Flip a bit in the last byte (inside the tag)
        let mut backing = buf.as_ref().to_vec();
        let last = backing.len() - 1;
        backing[last] ^= 0x01;
        let mut tmp = TmpBuf::new(&mut backing);
        assert!(
            cipher
                .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
                .is_err()
        );
    }

    /// Corrupted ciphertext body byte should cause decryption failure.
    #[test]
    fn ciphertext_corruption_detected() {
        let key = [0x42u8; 32];
        let nonce = Nonce([0x01u8; 12]);
        let plaintext = b"ciphertext corruption test";

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        let mut buf = Buf::new();
        buf.extend_from_slice(plaintext);
        cipher
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce)
            .unwrap();

        // Flip a bit in the first ciphertext byte
        let mut backing = buf.as_ref().to_vec();
        backing[0] ^= 0x01;
        let mut tmp = TmpBuf::new(&mut backing);
        assert!(
            cipher
                .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
                .is_err()
        );
    }

    /// Empty plaintext should produce tag-only output.
    #[test]
    fn empty_plaintext() {
        let key = [0x42u8; 32];
        let nonce = Nonce([0x01u8; 12]);

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        let mut buf = Buf::new();
        cipher
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce)
            .unwrap();

        // Should be exactly the tag length
        assert_eq!(buf.len(), CHACHA20_POLY1305_TAG_LEN);

        // Decrypt back to empty
        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        cipher
            .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
            .unwrap();
        assert_eq!(tmp.len(), 0);
    }

    /// Truncated ciphertext (shorter than tag) should fail.
    #[test]
    fn truncated_ciphertext_rejected() {
        let key = [0x42u8; 32];
        let nonce = Nonce([0x01u8; 12]);

        let mut cipher = ChaCha20Poly1305::new(&key).unwrap();

        let mut backing = vec![0u8; 8]; // less than 16 byte tag
        let mut tmp = TmpBuf::new(&mut backing);
        assert!(
            cipher
                .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
                .is_err()
        );
    }

    /// RFC 9001 Appendix A.5 ChaCha20 header protection test vector.
    #[test]
    fn encrypt_sn_rfc9001_vector() {
        let sn_key: [u8; 32] = [
            0x25, 0xa2, 0x82, 0xb9, 0xe8, 0x2f, 0x06, 0xf2, 0x1f, 0x48, 0x89, 0x17, 0xa4, 0xfc,
            0x8f, 0x1b, 0x73, 0x57, 0x36, 0x85, 0x60, 0x85, 0x97, 0xd0, 0xef, 0xcb, 0x07, 0x6b,
            0x0a, 0xb7, 0xa7, 0xa4,
        ];
        let sample: [u8; 16] = [
            0x5e, 0x5c, 0xd5, 0x5c, 0x41, 0xf6, 0x90, 0x80, 0x57, 0x5d, 0x79, 0x99, 0xc2, 0x5a,
            0x5b, 0xfb,
        ];

        let mask = TLS13_CHACHA20_POLY1305_SHA256.encrypt_sn(&sn_key, &sample);
        assert_eq!(&mask[..5], &[0xae, 0xfe, 0xfe, 0x7d, 0x03]);
    }

    /// Verify DTLS 1.3 suite metadata is consistent.
    #[test]
    fn suite_metadata() {
        assert_eq!(
            TLS13_CHACHA20_POLY1305_SHA256.suite(),
            Dtls13CipherSuite::CHACHA20_POLY1305_SHA256
        );
        assert_eq!(
            TLS13_CHACHA20_POLY1305_SHA256.hash_algorithm(),
            HashAlgorithm::SHA256
        );
        assert_eq!(TLS13_CHACHA20_POLY1305_SHA256.key_len(), 32);
        assert_eq!(TLS13_CHACHA20_POLY1305_SHA256.iv_len(), 12);
        assert_eq!(TLS13_CHACHA20_POLY1305_SHA256.tag_len(), 16);
    }

    /// SN encryption should be deterministic.
    #[test]
    fn encrypt_sn_deterministic() {
        let sn_key = [0x42u8; 32];
        let sample: [u8; 16] = [0x01u8; 16];
        let result = TLS13_CHACHA20_POLY1305_SHA256.encrypt_sn(&sn_key, &sample);
        let result_b = TLS13_CHACHA20_POLY1305_SHA256.encrypt_sn(&sn_key, &sample);
        assert_eq!(result, result_b);
    }

    /// Exercise create_cipher factory roundtrip.
    #[test]
    fn create_cipher_roundtrip() {
        let nonce = Nonce([0x01u8; 12]);
        let mut cipher = TLS13_CHACHA20_POLY1305_SHA256
            .create_cipher(&[0x42u8; 32])
            .unwrap();
        let mut buf = Buf::new();
        buf.extend_from_slice(b"dtls13 chacha");
        cipher
            .encrypt(&mut buf, Aad([0u8; 13].into()), nonce)
            .unwrap();
        let mut backing = buf.as_ref().to_vec();
        let mut tmp = TmpBuf::new(&mut backing);
        cipher
            .decrypt(&mut tmp, Aad([0u8; 13].into()), nonce)
            .unwrap();
        assert_eq!(tmp.as_ref(), b"dtls13 chacha");
    }
}
