//! ChaCha20-Poly1305 AEAD (RFC 8439).
//!
//! - ChaCha20 stream cipher: 256-bit key, 96-bit nonce, 32-bit counter
//! - Poly1305 one-time MAC: GF(2^130 - 5)
//! - AEAD construction: encrypt-then-MAC with padding and length encoding

// ---------------------------------------------------------------------------
// ChaCha20 quarter round (RFC 8439 §2.1)
// ---------------------------------------------------------------------------

fn quarter_round(state: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {
    state[a] = state[a].wrapping_add(state[b]);
    state[d] ^= state[a];
    state[d] = state[d].rotate_left(16);

    state[c] = state[c].wrapping_add(state[d]);
    state[b] ^= state[c];
    state[b] = state[b].rotate_left(12);

    state[a] = state[a].wrapping_add(state[b]);
    state[d] ^= state[a];
    state[d] = state[d].rotate_left(8);

    state[c] = state[c].wrapping_add(state[d]);
    state[b] ^= state[c];
    state[b] = state[b].rotate_left(7);
}

// ---------------------------------------------------------------------------
// ChaCha20 block function (RFC 8439 §2.3)
// ---------------------------------------------------------------------------

/// The ChaCha20 constants: "expand 32-byte k" as little-endian u32s.
const CONSTANTS: [u32; 4] = [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574];

fn chacha20_block(key: &[u8; 32], counter: u32, nonce: &[u8; 12]) -> [u8; 64] {
    let mut state = [0u32; 16];

    // Constants
    state[0] = CONSTANTS[0];
    state[1] = CONSTANTS[1];
    state[2] = CONSTANTS[2];
    state[3] = CONSTANTS[3];

    // Key (8 words, little-endian)
    for i in 0..8 {
        state[4 + i] = u32::from_le_bytes(key[4 * i..4 * i + 4].try_into().unwrap());
    }

    // Counter
    state[12] = counter;

    // Nonce (3 words, little-endian)
    for i in 0..3 {
        state[13 + i] = u32::from_le_bytes(nonce[4 * i..4 * i + 4].try_into().unwrap());
    }

    let initial = state;

    // 20 rounds (10 column rounds + 10 diagonal rounds)
    for _ in 0..10 {
        // Column rounds
        quarter_round(&mut state, 0, 4, 8, 12);
        quarter_round(&mut state, 1, 5, 9, 13);
        quarter_round(&mut state, 2, 6, 10, 14);
        quarter_round(&mut state, 3, 7, 11, 15);
        // Diagonal rounds
        quarter_round(&mut state, 0, 5, 10, 15);
        quarter_round(&mut state, 1, 6, 11, 12);
        quarter_round(&mut state, 2, 7, 8, 13);
        quarter_round(&mut state, 3, 4, 9, 14);
    }

    // Add initial state
    for i in 0..16 {
        state[i] = state[i].wrapping_add(initial[i]);
    }

    // Serialize to bytes (little-endian)
    let mut out = [0u8; 64];
    for i in 0..16 {
        out[4 * i..4 * i + 4].copy_from_slice(&state[i].to_le_bytes());
    }
    out
}

// ---------------------------------------------------------------------------
// ChaCha20 encryption (RFC 8439 §2.4)
// ---------------------------------------------------------------------------

fn chacha20_encrypt(key: &[u8; 32], counter: u32, nonce: &[u8; 12], data: &[u8]) -> Vec<u8> {
    let mut result = Vec::with_capacity(data.len());
    let mut block_counter = counter;
    let mut offset = 0;

    while offset < data.len() {
        let keystream = chacha20_block(key, block_counter, nonce);
        let remaining = (data.len() - offset).min(64);
        for i in 0..remaining {
            result.push(data[offset + i] ^ keystream[i]);
        }
        offset += remaining;
        block_counter = block_counter.wrapping_add(1);
    }

    result
}

// ---------------------------------------------------------------------------
// Poly1305 MAC (RFC 8439 §2.5)
// ---------------------------------------------------------------------------

/// Poly1305 one-time authenticator.
///
/// Computes the MAC over `data` using the one-time key `key` (32 bytes).
/// The key is split into `r` (clamped, first 16 bytes) and `s` (last 16 bytes).
///
/// All arithmetic is in GF(2^130 - 5) using 5 limbs of 26 bits each.
fn poly1305_mac(key: &[u8; 32], data: &[u8]) -> [u8; 16] {
    // Clamp r per RFC 8439 §2.5
    let mut r_bytes = [0u8; 16];
    r_bytes.copy_from_slice(&key[0..16]);
    r_bytes[3] &= 15;
    r_bytes[7] &= 15;
    r_bytes[11] &= 15;
    r_bytes[15] &= 15;
    r_bytes[4] &= 252;
    r_bytes[8] &= 252;
    r_bytes[12] &= 252;

    // r as 5 limbs of 26 bits
    let r0 = (u32::from_le_bytes(r_bytes[0..4].try_into().unwrap())) & 0x3ffffff;
    let r1 = (u32::from_le_bytes(r_bytes[3..7].try_into().unwrap()) >> 2) & 0x3ffffff;
    let r2 = (u32::from_le_bytes(r_bytes[6..10].try_into().unwrap()) >> 4) & 0x3ffffff;
    let r3 = (u32::from_le_bytes(r_bytes[9..13].try_into().unwrap()) >> 6) & 0x3ffffff;
    let r4 = (u32::from_le_bytes(r_bytes[12..16].try_into().unwrap()) >> 8) & 0x3ffffff;

    // s (second half of key)
    let s = &key[16..32];

    // Precompute 5*r[i] for reduction
    let s1 = r1 * 5;
    let s2 = r2 * 5;
    let s3 = r3 * 5;
    let s4 = r4 * 5;

    // Accumulator h = 0
    let mut h0: u32 = 0;
    let mut h1: u32 = 0;
    let mut h2: u32 = 0;
    let mut h3: u32 = 0;
    let mut h4: u32 = 0;

    // Process message in 16-byte blocks
    let mut offset = 0;
    while offset < data.len() {
        let remaining = data.len() - offset;
        let block_len = remaining.min(16);

        // Read block and add padding byte
        let mut block = [0u8; 17];
        block[..block_len].copy_from_slice(&data[offset..offset + block_len]);
        block[block_len] = 1; // hibit

        // Convert to 5 limbs of 26 bits
        let t0 = u32::from_le_bytes(block[0..4].try_into().unwrap());
        let t1 = u32::from_le_bytes(block[4..8].try_into().unwrap());
        let t2 = u32::from_le_bytes(block[8..12].try_into().unwrap());
        let t3 = u32::from_le_bytes(block[12..16].try_into().unwrap());
        let hibit = block[16] as u32;

        h0 = h0.wrapping_add(t0 & 0x3ffffff);
        h1 = h1.wrapping_add(((t0 >> 26) | (t1 << 6)) & 0x3ffffff);
        h2 = h2.wrapping_add(((t1 >> 20) | (t2 << 12)) & 0x3ffffff);
        h3 = h3.wrapping_add(((t2 >> 14) | (t3 << 18)) & 0x3ffffff);
        h4 = h4.wrapping_add((t3 >> 8) | (hibit << 24));

        // h *= r (mod 2^130 - 5)
        let d0 = (h0 as u64) * (r0 as u64)
            + (h1 as u64) * (s4 as u64)
            + (h2 as u64) * (s3 as u64)
            + (h3 as u64) * (s2 as u64)
            + (h4 as u64) * (s1 as u64);
        let d1 = (h0 as u64) * (r1 as u64)
            + (h1 as u64) * (r0 as u64)
            + (h2 as u64) * (s4 as u64)
            + (h3 as u64) * (s3 as u64)
            + (h4 as u64) * (s2 as u64);
        let d2 = (h0 as u64) * (r2 as u64)
            + (h1 as u64) * (r1 as u64)
            + (h2 as u64) * (r0 as u64)
            + (h3 as u64) * (s4 as u64)
            + (h4 as u64) * (s3 as u64);
        let d3 = (h0 as u64) * (r3 as u64)
            + (h1 as u64) * (r2 as u64)
            + (h2 as u64) * (r1 as u64)
            + (h3 as u64) * (r0 as u64)
            + (h4 as u64) * (s4 as u64);
        let d4 = (h0 as u64) * (r4 as u64)
            + (h1 as u64) * (r3 as u64)
            + (h2 as u64) * (r2 as u64)
            + (h3 as u64) * (r1 as u64)
            + (h4 as u64) * (r0 as u64);

        // Partial reduction mod 2^130 - 5
        let mut c: u32;
        c = (d0 >> 26) as u32;
        h0 = d0 as u32 & 0x3ffffff;
        let d1 = d1 + c as u64;
        c = (d1 >> 26) as u32;
        h1 = d1 as u32 & 0x3ffffff;
        let d2 = d2 + c as u64;
        c = (d2 >> 26) as u32;
        h2 = d2 as u32 & 0x3ffffff;
        let d3 = d3 + c as u64;
        c = (d3 >> 26) as u32;
        h3 = d3 as u32 & 0x3ffffff;
        let d4 = d4 + c as u64;
        c = (d4 >> 26) as u32;
        h4 = d4 as u32 & 0x3ffffff;
        h0 = h0.wrapping_add(c * 5);
        c = h0 >> 26;
        h0 &= 0x3ffffff;
        h1 = h1.wrapping_add(c);

        offset += block_len;
    }

    // Final reduction: fully carry h
    let mut c = h1 >> 26;
    h1 &= 0x3ffffff;
    h2 = h2.wrapping_add(c);
    c = h2 >> 26;
    h2 &= 0x3ffffff;
    h3 = h3.wrapping_add(c);
    c = h3 >> 26;
    h3 &= 0x3ffffff;
    h4 = h4.wrapping_add(c);
    c = h4 >> 26;
    h4 &= 0x3ffffff;
    h0 = h0.wrapping_add(c * 5);
    c = h0 >> 26;
    h0 &= 0x3ffffff;
    h1 = h1.wrapping_add(c);

    // Compute h + -(2^130 - 5) = h - p
    let mut g0 = h0.wrapping_add(5);
    c = g0 >> 26;
    g0 &= 0x3ffffff;
    let mut g1 = h1.wrapping_add(c);
    c = g1 >> 26;
    g1 &= 0x3ffffff;
    let mut g2 = h2.wrapping_add(c);
    c = g2 >> 26;
    g2 &= 0x3ffffff;
    let mut g3 = h3.wrapping_add(c);
    c = g3 >> 26;
    g3 &= 0x3ffffff;
    let g4 = h4.wrapping_add(c).wrapping_sub(1 << 26);

    // Select h or g based on whether g overflowed (constant-time)
    let mask = (g4 >> 31).wrapping_sub(1); // 0xffffffff if g4 >= 0, 0 if g4 < 0
    g0 &= mask;
    g1 &= mask;
    g2 &= mask;
    g3 &= mask;
    let g4 = g4 & mask;
    let nmask = !mask;
    h0 = (h0 & nmask) | g0;
    h1 = (h1 & nmask) | g1;
    h2 = (h2 & nmask) | g2;
    h3 = (h3 & nmask) | g3;
    h4 = (h4 & nmask) | g4;

    // Reassemble h as a 128-bit number (mod 2^128) and add s
    let h_val: u128 = (h0 as u128)
        | ((h1 as u128) << 26)
        | ((h2 as u128) << 52)
        | ((h3 as u128) << 78)
        | ((h4 as u128) << 104);
    let s_val = u128::from_le_bytes(s.try_into().unwrap());
    let result = h_val.wrapping_add(s_val);
    result.to_le_bytes()
}

// ---------------------------------------------------------------------------
// Constant-time tag comparison
// ---------------------------------------------------------------------------

fn ct_eq_16(a: &[u8; 16], b: &[u8; 16]) -> bool {
    let mut diff = 0u8;
    for i in 0..16 {
        diff |= a[i] ^ b[i];
    }
    diff == 0
}

// ---------------------------------------------------------------------------
// ChaCha20-Poly1305 AEAD (RFC 8439 §2.8)
// ---------------------------------------------------------------------------

/// Construct the Poly1305 input per RFC 8439 §2.8:
/// AAD ∥ pad(AAD) ∥ ciphertext ∥ pad(ciphertext) ∥ len(AAD) ∥ len(CT)
fn build_mac_data(aad: &[u8], ciphertext: &[u8]) -> Vec<u8> {
    let mut mac_data = Vec::new();

    // AAD + padding to 16-byte boundary
    mac_data.extend_from_slice(aad);
    let pad_aad = (16 - (aad.len() % 16)) % 16;
    mac_data.extend(std::iter::repeat_n(0u8, pad_aad));

    // Ciphertext + padding to 16-byte boundary
    mac_data.extend_from_slice(ciphertext);
    let pad_ct = (16 - (ciphertext.len() % 16)) % 16;
    mac_data.extend(std::iter::repeat_n(0u8, pad_ct));

    // Lengths as 64-bit little-endian
    mac_data.extend_from_slice(&(aad.len() as u64).to_le_bytes());
    mac_data.extend_from_slice(&(ciphertext.len() as u64).to_le_bytes());

    mac_data
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// ChaCha20-Poly1305 AEAD encrypt (RFC 8439).
///
/// Returns `(ciphertext, 128-bit tag)`.
pub fn chacha20_poly1305_encrypt(
    key: &[u8; 32],
    nonce: &[u8; 12],
    plaintext: &[u8],
    aad: &[u8],
) -> (Vec<u8>, [u8; 16]) {
    // Generate Poly1305 one-time key from first ChaCha20 block (counter = 0)
    let otk_block = chacha20_block(key, 0, nonce);
    let mut otk = [0u8; 32];
    otk.copy_from_slice(&otk_block[0..32]);

    // Encrypt plaintext with ChaCha20 (counter starts at 1)
    let ciphertext = chacha20_encrypt(key, 1, nonce, plaintext);

    // Compute tag
    let mac_data = build_mac_data(aad, &ciphertext);
    let tag = poly1305_mac(&otk, &mac_data);

    (ciphertext, tag)
}

/// ChaCha20-Poly1305 AEAD decrypt (RFC 8439).
///
/// Returns `None` if the authentication tag is invalid.
pub fn chacha20_poly1305_decrypt(
    key: &[u8; 32],
    nonce: &[u8; 12],
    ciphertext: &[u8],
    aad: &[u8],
    tag: &[u8; 16],
) -> Option<Vec<u8>> {
    // Generate Poly1305 one-time key
    let otk_block = chacha20_block(key, 0, nonce);
    let mut otk = [0u8; 32];
    otk.copy_from_slice(&otk_block[0..32]);

    // Verify tag before decrypting
    let mac_data = build_mac_data(aad, ciphertext);
    let expected_tag = poly1305_mac(&otk, &mac_data);

    if !ct_eq_16(&expected_tag, tag) {
        return None;
    }

    // Decrypt ciphertext with ChaCha20 (counter starts at 1)
    let plaintext = chacha20_encrypt(key, 1, nonce, ciphertext);
    Some(plaintext)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn from_hex(s: &str) -> Vec<u8> {
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    fn hex(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{b:02x}")).collect()
    }

    // -----------------------------------------------------------------------
    // RFC 8439 §2.1.1 — Quarter Round Test Vector
    // -----------------------------------------------------------------------

    #[test]
    fn quarter_round_test_vector() {
        let mut state = [0u32; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        // Apply quarter round to indices 0,1,2,3 (need a 16-element state)
        quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!(state[0], 0xea2a92f4);
        assert_eq!(state[1], 0xcb1cf8ce);
        assert_eq!(state[2], 0x4581472e);
        assert_eq!(state[3], 0x5881c4bb);
    }

    // -----------------------------------------------------------------------
    // RFC 8439 §2.3.2 — ChaCha20 Block Function Test Vector
    // -----------------------------------------------------------------------

    #[test]
    fn chacha20_block_test_vector() {
        let key: [u8; 32] = from_hex(
            "000102030405060708090a0b0c0d0e0f\
             101112131415161718191a1b1c1d1e1f",
        )
        .try_into()
        .unwrap();
        let nonce: [u8; 12] = from_hex("000000090000004a00000000").try_into().unwrap();
        let counter: u32 = 1;

        let block = chacha20_block(&key, counter, &nonce);

        let expected = from_hex(
            "10f1e7e4d13b5915500fdd1fa32071c4\
             c7d1f4c733c068030422aa9ac3d46c4e\
             d2826446079faa0914c2d705d98b02a2\
             b5129cd1de164eb9cbd083e8a2503c4e",
        );
        assert_eq!(block.to_vec(), expected);
    }

    // -----------------------------------------------------------------------
    // RFC 8439 §2.4.2 — ChaCha20 Encryption Test Vector
    // -----------------------------------------------------------------------

    #[test]
    fn chacha20_encryption_test_vector() {
        let key: [u8; 32] = from_hex(
            "000102030405060708090a0b0c0d0e0f\
             101112131415161718191a1b1c1d1e1f",
        )
        .try_into()
        .unwrap();
        let nonce: [u8; 12] = from_hex("000000000000004a00000000").try_into().unwrap();
        let counter: u32 = 1;

        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 ciphertext = chacha20_encrypt(&key, counter, &nonce, plaintext);

        let expected = from_hex(
            "6e2e359a2568f98041ba0728dd0d6981\
             e97e7aec1d4360c20a27afccfd9fae0b\
             f91b65c5524733ab8f593dabcd62b357\
             1639d624e65152ab8f530c359f0861d8\
             07ca0dbf500d6a6156a38e088a22b65e\
             52bc514d16ccf806818ce91ab7793736\
             5af90bbf74a35be6b40b8eedf2785e42\
             874d",
        );
        assert_eq!(hex(&ciphertext), hex(&expected));
    }

    // -----------------------------------------------------------------------
    // RFC 8439 §2.5.2 — Poly1305 MAC Test Vector
    // -----------------------------------------------------------------------

    #[test]
    fn poly1305_mac_test_vector() {
        let key: [u8; 32] = from_hex(
            "85d6be7857556d337f4452fe42d506a8\
             0103808afb0db2fd4abff6af4149f51b",
        )
        .try_into()
        .unwrap();
        let msg = b"Cryptographic Forum Research Group";
        let tag = poly1305_mac(&key, msg);
        assert_eq!(hex(&tag), "a8061dc1305136c6c22b8baf0c0127a9");
    }

    // -----------------------------------------------------------------------
    // RFC 8439 §2.6.2 — Poly1305 Key Generation Test Vector
    // -----------------------------------------------------------------------

    #[test]
    fn poly1305_key_generation_test_vector() {
        let key: [u8; 32] = from_hex(
            "808182838485868788898a8b8c8d8e8f\
             909192939495969798999a9b9c9d9e9f",
        )
        .try_into()
        .unwrap();
        let nonce: [u8; 12] = from_hex("000000000001020304050607").try_into().unwrap();

        let otk_block = chacha20_block(&key, 0, &nonce);
        let otk = &otk_block[0..32];

        assert_eq!(
            hex(otk),
            "8ad5a08b905f81cc815040274ab29471\
             a833b637e3fd0da508dbb8e2fdd1a646"
        );
    }

    // -----------------------------------------------------------------------
    // RFC 8439 §2.8.2 — AEAD Test Vector
    // -----------------------------------------------------------------------

    #[test]
    fn aead_encrypt_test_vector() {
        let key: [u8; 32] = from_hex(
            "808182838485868788898a8b8c8d8e8f\
             909192939495969798999a9b9c9d9e9f",
        )
        .try_into()
        .unwrap();
        let nonce: [u8; 12] = from_hex("070000004041424344454647").try_into().unwrap();
        let aad = from_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) = chacha20_poly1305_encrypt(&key, &nonce, plaintext, &aad);

        let expected_ct = from_hex(
            "d31a8d34648e60db7b86afbc53ef7ec2\
             a4aded51296e08fea9e2b5a736ee62d6\
             3dbea45e8ca9671282fafb69da92728b\
             1a71de0a9e060b2905d6a5b67ecd3b36\
             92ddbd7f2d778b8c9803aee328091b58\
             fab324e4fad675945585808b4831d7bc\
             3ff4def08e4b7a9de576d26586cec64b\
             6116",
        );
        assert_eq!(hex(&ct), hex(&expected_ct));
        assert_eq!(hex(&tag), "1ae10b594f09e26a7e902ecbd0600691");
    }

    #[test]
    fn aead_decrypt_test_vector() {
        let key: [u8; 32] = from_hex(
            "808182838485868788898a8b8c8d8e8f\
             909192939495969798999a9b9c9d9e9f",
        )
        .try_into()
        .unwrap();
        let nonce: [u8; 12] = from_hex("070000004041424344454647").try_into().unwrap();
        let aad = from_hex("50515253c0c1c2c3c4c5c6c7");
        let ct = from_hex(
            "d31a8d34648e60db7b86afbc53ef7ec2\
             a4aded51296e08fea9e2b5a736ee62d6\
             3dbea45e8ca9671282fafb69da92728b\
             1a71de0a9e060b2905d6a5b67ecd3b36\
             92ddbd7f2d778b8c9803aee328091b58\
             fab324e4fad675945585808b4831d7bc\
             3ff4def08e4b7a9de576d26586cec64b\
             6116",
        );
        let tag: [u8; 16] = from_hex("1ae10b594f09e26a7e902ecbd0600691")
            .try_into()
            .unwrap();

        let pt = chacha20_poly1305_decrypt(&key, &nonce, &ct, &aad, &tag).unwrap();
        assert_eq!(
            pt,
            b"Ladies and Gentlemen of the class of '99: \
If I could offer you only one tip for the future, sunscreen would be it."
                .to_vec()
        );
    }

    // -----------------------------------------------------------------------
    // Tag tamper detection
    // -----------------------------------------------------------------------

    #[test]
    fn decrypt_rejects_tampered_tag() {
        let key = [0x42u8; 32];
        let nonce = [0u8; 12];
        let (ct, mut tag) = chacha20_poly1305_encrypt(&key, &nonce, b"hello", &[]);
        tag[0] ^= 1;
        assert!(chacha20_poly1305_decrypt(&key, &nonce, &ct, &[], &tag).is_none());
    }

    #[test]
    fn decrypt_rejects_tampered_ciphertext() {
        let key = [0x42u8; 32];
        let nonce = [0u8; 12];
        let (mut ct, tag) = chacha20_poly1305_encrypt(&key, &nonce, b"hello", &[]);
        ct[0] ^= 1;
        assert!(chacha20_poly1305_decrypt(&key, &nonce, &ct, &[], &tag).is_none());
    }

    #[test]
    fn decrypt_rejects_tampered_aad() {
        let key = [0x42u8; 32];
        let nonce = [0u8; 12];
        let aad = b"metadata";
        let (ct, tag) = chacha20_poly1305_encrypt(&key, &nonce, b"hello", aad);
        let bad_aad = b"Metadata";
        assert!(chacha20_poly1305_decrypt(&key, &nonce, &ct, bad_aad, &tag).is_none());
    }

    // -----------------------------------------------------------------------
    // Round-trip
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_empty_plaintext() {
        let key = [0xabu8; 32];
        let nonce = [0x01u8; 12];
        let (ct, tag) = chacha20_poly1305_encrypt(&key, &nonce, &[], &[]);
        assert!(ct.is_empty());
        let pt = chacha20_poly1305_decrypt(&key, &nonce, &ct, &[], &tag).unwrap();
        assert!(pt.is_empty());
    }

    #[test]
    fn roundtrip_with_aad() {
        let key: [u8; 32] = from_hex(
            "deadbeefdeadbeefdeadbeefdeadbeef\
             deadbeefdeadbeefdeadbeefdeadbeef",
        )
        .try_into()
        .unwrap();
        let nonce: [u8; 12] = from_hex("0102030405060708090a0b0c").try_into().unwrap();
        let pt = b"The quick brown fox jumps over the lazy dog";
        let aad = b"additional authenticated data";

        let (ct, tag) = chacha20_poly1305_encrypt(&key, &nonce, pt, aad);
        let recovered = chacha20_poly1305_decrypt(&key, &nonce, &ct, aad, &tag).unwrap();
        assert_eq!(recovered, pt.to_vec());
    }

    #[test]
    fn roundtrip_large_message() {
        let key = [0xffu8; 32];
        let nonce = [0u8; 12];
        let pt: Vec<u8> = (0..1024).map(|i| (i % 256) as u8).collect();
        let aad = b"large message test";

        let (ct, tag) = chacha20_poly1305_encrypt(&key, &nonce, &pt, aad);
        let recovered = chacha20_poly1305_decrypt(&key, &nonce, &ct, aad, &tag).unwrap();
        assert_eq!(recovered, pt);
    }
}