//! ChaCha20-Poly1305 AEAD (RFC 8439).
//!
//! Authenticated encryption with associated data, built from
//! ChaCha20 (the stream cipher) and Poly1305 (the one-time MAC).
//! This is the AEAD used by TLS 1.3, Noise, Signal, WireGuard,
//! QUIC, OpenSSH, and most modern protocols where AES-GCM is not
//! the right choice -- typically because the platform lacks
//! AES-NI and a constant-time AES would be too slow.
//!
//! # Construction (RFC 8439 §2.8)
//!
//! 1. Initialise ChaCha20 with `(key, nonce, counter = 0)` and
//!    take the first 32 bytes of the resulting keystream as the
//!    one-time Poly1305 key.
//! 2. Re-initialise ChaCha20 with `(key, nonce, counter = 1)` and
//!    XOR the plaintext with the keystream to produce the
//!    ciphertext.
//! 3. Compute the Poly1305 tag over:
//!
//! ```text
//!     AAD                                ‖ pad16(AAD)
//!     ciphertext                         ‖ pad16(ciphertext)
//!     len(AAD) as u64 little-endian
//!     len(ciphertext) as u64 little-endian
//! ```
//!
//! Where `pad16(x)` is `0..(16 - len(x) mod 16) mod 16` zero bytes.
//!
//! 4. Output `(ciphertext, tag)`. The tag is 16 bytes.
//!
//! Decryption reverses the process and verifies the tag in
//! constant time before returning the plaintext.
//!
//! # API
//!
//! ```rust,ignore
//! use arcana::cipher::chacha20poly1305::ChaCha20Poly1305;
//!
//! let key:   [u8; 32] = /* shared secret */;
//! let nonce: [u8; 12] = /* MUST be unique per (key, message) */;
//! let aad = b"associated header";
//! let pt  = b"top secret message";
//!
//! let (ct, tag) = ChaCha20Poly1305::encrypt(&key, &nonce, aad, pt);
//! let pt_back = ChaCha20Poly1305::decrypt(&key, &nonce, aad, &ct, &tag)
//!                 .expect("authentic");
//! assert_eq!(pt_back, pt);
//! ```
//!
//! # Nonce reuse warning
//!
//! Reusing a nonce with the same key on a different `(plaintext,
//! aad)` is **catastrophic**: it leaks the XOR of the two plaintexts
//! AND lets the attacker forge arbitrary messages under that key.
//! Always draw a fresh random 12-byte nonce or use a counter that
//! is guaranteed unique for the lifetime of the key.
//!
//! # Tests
//!
//! Pinned against RFC 8439 §2.8.2 (the canonical "Ladies and
//! Gentlemen" AEAD test vector) byte-for-byte.

use super::chacha20::ChaCha20;
use super::poly1305::Poly1305;

// ============================================================================
// Public API
// ============================================================================

/// ChaCha20-Poly1305 AEAD per RFC 8439.
///
/// Stateless tag struct -- the per-message state lives in the
/// ChaCha20 / Poly1305 instances. Exposed as a unit struct so the
/// API matches the AES-GCM `Gcm` struct in `cipher::modes`.
pub struct ChaCha20Poly1305;

impl ChaCha20Poly1305 {
    /// Encrypt and authenticate a message.
    ///
    /// Returns `(ciphertext, tag)` where `ciphertext.len() ==
    /// plaintext.len()` and `tag` is exactly 16 bytes. Both must
    /// be transmitted to the receiver alongside the nonce and AAD.
    pub fn encrypt(key: &[u8; 32], nonce: &[u8; 12], aad: &[u8], plaintext: &[u8]) -> (Vec<u8>, [u8; 16]) {
        // Step 1: derive the one-time Poly1305 key from ChaCha20
        // block 0 (counter = 0).
        let poly_key = poly_key_gen(key, nonce);

        // Step 2: encrypt with ChaCha20 starting at counter = 1.
        let mut ct = plaintext.to_vec();
        let mut cipher = ChaCha20::new(key, nonce, 1);
        cipher.apply_keystream(&mut ct);

        // Step 3: compute the Poly1305 tag over the AEAD layout.
        let tag = compute_tag(&poly_key, aad, &ct);

        (ct, tag)
    }

    /// Decrypt and verify a ciphertext.
    ///
    /// Returns `Some(plaintext)` only if `tag` is the correct MAC
    /// for `(aad, ciphertext)` under `(key, nonce)`. The tag is
    /// compared in constant time -- the function execution time
    /// does not leak which byte of the tag was wrong.
    ///
    /// Returns `None` if the tag does not verify. **Callers MUST
    /// NOT use the returned plaintext if `None` is returned**, and
    /// in particular must not log it, hash it, or branch on its
    /// contents -- the only correct response to a bad tag is to
    /// abort the protocol.
    pub fn decrypt(key: &[u8; 32], nonce: &[u8; 12], aad: &[u8], ciphertext: &[u8], tag: &[u8; 16]) -> Option<Vec<u8>> {
        // Recompute the expected tag and compare in constant time.
        let poly_key = poly_key_gen(key, nonce);
        let expected_tag = compute_tag(&poly_key, aad, ciphertext);

        let mut diff = 0u8;
        for i in 0..16 {
            diff |= expected_tag[i] ^ tag[i];
        }
        if diff != 0 {
            return None;
        }

        // Tag is good -- decrypt and return.
        let mut pt = ciphertext.to_vec();
        let mut cipher = ChaCha20::new(key, nonce, 1);
        cipher.apply_keystream(&mut pt);
        Some(pt)
    }
}

// ============================================================================
// Internal helpers
// ============================================================================

/// RFC 8439 §2.6 `poly1305_key_gen`: take the first 32 bytes of
/// `ChaCha20(key, nonce, counter = 0)` as the one-time Poly1305 key.
fn poly_key_gen(key: &[u8; 32], nonce: &[u8; 12]) -> [u8; 32] {
    // Apply 32 zero bytes through the keystream to extract the
    // first 32 bytes of block 0.
    let mut buf = [0u8; 32];
    let mut cipher = ChaCha20::new(key, nonce, 0);
    cipher.apply_keystream(&mut buf);
    buf
}

/// Compute the Poly1305 tag over the AEAD layout (RFC 8439 §2.8):
///
/// ```text
///   AAD || pad16(AAD)
///   CT  || pad16(CT)
///   len(AAD) as u64 LE
///   len(CT)  as u64 LE
/// ```
///
/// Returns the 16-byte tag.
fn compute_tag(poly_key: &[u8; 32], aad: &[u8], ct: &[u8]) -> [u8; 16] {
    let mut poly = Poly1305::new(poly_key);
    poly.update(aad);
    poly.update(&zero_pad_to_16(aad.len()));
    poly.update(ct);
    poly.update(&zero_pad_to_16(ct.len()));
    let mut len_bytes = [0u8; 16];
    len_bytes[0..8].copy_from_slice(&(aad.len() as u64).to_le_bytes());
    len_bytes[8..16].copy_from_slice(&(ct.len() as u64).to_le_bytes());
    poly.update(&len_bytes);
    poly.finalize()
}

/// Returns the zero pad needed to round `n` up to a multiple of 16.
/// Length is in `0..16`.
fn zero_pad_to_16(n: usize) -> Vec<u8> {
    let pad_len = (16 - (n % 16)) % 16;
    vec![0u8; pad_len]
}

// ============================================================================
// Tests (RFC 8439 pinned vectors)
// ============================================================================

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

    fn hex(s: &str) -> Vec<u8> {
        let s: String = s.chars().filter(|c| !c.is_whitespace()).collect();
        assert!(s.len() % 2 == 0, "odd hex length: {}", s.len());
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    fn hex_arr<const N: usize>(s: &str) -> [u8; N] {
        let v = hex(s);
        assert_eq!(v.len(), N);
        let mut out = [0u8; N];
        out.copy_from_slice(&v);
        out
    }

    /// RFC 8439 §2.6.2 `poly1305_key_gen` test vector.
    ///
    /// Key:   80:81:82:..:9f
    /// Nonce: 00:00:00:00:00:01:02:03:04:05:06:07
    /// Out:   8ad5a08b905f81cc815040274ab29471
    ///        a833b637e3fd0da508dbb8e2fdd1a646
    #[test]
    fn rfc8439_2_6_2_poly_key_gen() {
        let key: [u8; 32] = hex_arr("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f");
        let nonce: [u8; 12] = hex_arr("000000000001020304050607");
        let pk = poly_key_gen(&key, &nonce);
        let expected = hex("8ad5a08b905f81cc815040274ab29471\
             a833b637e3fd0da508dbb8e2fdd1a646");
        assert_eq!(pk.to_vec(), expected);
    }

    /// RFC 8439 §2.8.2 full AEAD test vector ("Ladies and Gentlemen").
    ///
    /// Key:   80:81:82:..:9f
    /// Nonce: 07:00:00:00:40:41:42:43:44:45:46:47
    /// AAD:   50:51:52:53:c0:c1:c2:c3:c4:c5:c6:c7
    /// Plaintext: "Ladies and Gentlemen of the class of '99: ..."
    /// Expected ciphertext + tag pinned below.
    #[test]
    fn rfc8439_2_8_2_aead_vector() {
        let key: [u8; 32] = hex_arr("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f");
        let nonce: [u8; 12] = hex_arr("070000004041424344454647");
        let aad = 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 (ct, tag) = ChaCha20Poly1305::encrypt(&key, &nonce, &aad, plaintext);

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

        assert_eq!(ct, expected_ct);
        assert_eq!(tag.to_vec(), expected_tag);

        // Decrypt round-trip with the pinned tag.
        let mut tag_arr = [0u8; 16];
        tag_arr.copy_from_slice(&expected_tag);
        let pt = ChaCha20Poly1305::decrypt(&key, &nonce, &aad, &expected_ct, &tag_arr).expect("authentic");
        assert_eq!(pt, plaintext);
    }

    /// Encrypt then decrypt round-trip on an arbitrary message.
    #[test]
    fn aead_roundtrip_random_inputs() {
        let key = [0x42u8; 32];
        let nonce = [0xa5u8; 12];
        let aad = b"some context";
        let pt = b"hello world; this is a test of moderate length to span more than one ChaCha20 block.";

        let (ct, tag) = ChaCha20Poly1305::encrypt(&key, &nonce, aad, pt);
        assert_ne!(ct.as_slice(), pt.as_slice());
        let back = ChaCha20Poly1305::decrypt(&key, &nonce, aad, &ct, &tag).expect("authentic");
        assert_eq!(back.as_slice(), pt.as_slice());
    }

    /// Decrypt must reject a tampered ciphertext byte.
    #[test]
    fn aead_rejects_tampered_ciphertext() {
        let key = [0x01u8; 32];
        let nonce = [0x02u8; 12];
        let pt = b"do not modify";

        let (mut ct, tag) = ChaCha20Poly1305::encrypt(&key, &nonce, b"", pt);
        ct[0] ^= 0x01;
        assert!(ChaCha20Poly1305::decrypt(&key, &nonce, b"", &ct, &tag).is_none());
    }

    /// Decrypt must reject a tampered tag byte.
    #[test]
    fn aead_rejects_tampered_tag() {
        let key = [0x01u8; 32];
        let nonce = [0x02u8; 12];
        let pt = b"do not modify";

        let (ct, mut tag) = ChaCha20Poly1305::encrypt(&key, &nonce, b"", pt);
        tag[0] ^= 0x01;
        assert!(ChaCha20Poly1305::decrypt(&key, &nonce, b"", &ct, &tag).is_none());
    }

    /// Decrypt must reject if the AAD changes -- AAD is part of the
    /// MAC input. This is the property that makes "associated data"
    /// authenticated even though it is not encrypted.
    #[test]
    fn aead_rejects_modified_aad() {
        let key = [0xffu8; 32];
        let nonce = [0x10u8; 12];
        let aad = b"context-A";
        let pt = b"shared payload";

        let (ct, tag) = ChaCha20Poly1305::encrypt(&key, &nonce, aad, pt);
        assert!(ChaCha20Poly1305::decrypt(&key, &nonce, b"context-B", &ct, &tag).is_none());
    }

    /// Decrypt must reject if the wrong key is used.
    #[test]
    fn aead_rejects_wrong_key() {
        let key1 = [0x33u8; 32];
        let mut key2 = key1;
        key2[0] ^= 0x01;
        let nonce = [0x44u8; 12];
        let pt = b"sensitive";

        let (ct, tag) = ChaCha20Poly1305::encrypt(&key1, &nonce, b"", pt);
        assert!(ChaCha20Poly1305::decrypt(&key2, &nonce, b"", &ct, &tag).is_none());
    }

    /// Empty plaintext is a valid input -- the ciphertext is also
    /// empty and the tag still authenticates the AAD and the
    /// length fields.
    #[test]
    fn aead_empty_plaintext() {
        let key = [0x55u8; 32];
        let nonce = [0x66u8; 12];
        let aad = b"only-context";
        let (ct, tag) = ChaCha20Poly1305::encrypt(&key, &nonce, aad, b"");
        assert!(ct.is_empty());
        let back = ChaCha20Poly1305::decrypt(&key, &nonce, aad, &ct, &tag).unwrap();
        assert!(back.is_empty());

        // Modifying the AAD must still be detected on empty plaintext.
        assert!(ChaCha20Poly1305::decrypt(&key, &nonce, b"other", &ct, &tag).is_none());
    }
}
