//! XChaCha20-Poly1305 AEAD (draft-irtf-cfrg-xchacha).
//!
//! Extension of ChaCha20-Poly1305 (RFC 8439) to a **24-byte nonce**
//! via the HChaCha20 subkey derivation. The larger nonce makes it
//! safe to pick nonces randomly without tracking a counter — the
//! birthday bound becomes `2^96` instead of `2^48` for the 12-byte
//! IETF nonce.
//!
//! # Construction
//!
//! Given a 32-byte key `K` and a 24-byte nonce `N`:
//!
//! 1. Split `N` into `N[0..16]` (for HChaCha20) and `N[16..24]`.
//! 2. `subkey = HChaCha20(K, N[0..16])` — a 32-byte derived key.
//! 3. `nonce' = 0x00000000 || N[16..24]` — a 12-byte IETF nonce.
//! 4. Run `ChaCha20-Poly1305(subkey, nonce', aad, plaintext)`.
//!
//! Used by libsodium (`crypto_aead_xchacha20poly1305_ietf_*`),
//! Signal, Age, WireGuard handshake, and many modern protocols
//! that want random nonces without the 2^48 cap.

use crate::cipher::chacha20::quarter_round;
use crate::cipher::chacha20poly1305::ChaCha20Poly1305;

// ============================================================================
// HChaCha20 (draft-irtf-cfrg-xchacha §2.2)
// ============================================================================

/// HChaCha20 subkey derivation: 32-byte key + 16-byte input → 32-byte output.
///
/// Unlike ChaCha20 block, HChaCha20 does **not** add the initial
/// state to the output: it serializes the post-round state directly.
fn hchacha20(key: &[u8; 32], input: &[u8; 16]) -> [u8; 32] {
    let mut state = [0u32; 16];

    // Constants: "expand 32-byte k" as four little-endian u32.
    state[0] = 0x6170_7865;
    state[1] = 0x3320_646e;
    state[2] = 0x7962_2d32;
    state[3] = 0x6b20_6574;

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

    // 16-byte input occupies positions [12..16] (where ChaCha20 has
    // counter + 12-byte nonce).
    for i in 0..4 {
        state[12 + i] = u32::from_le_bytes(input[4 * i..4 * i + 4].try_into().unwrap());
    }

    // 20 rounds = 10 double-rounds (same as ChaCha20).
    for _ in 0..10 {
        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);
        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);
    }

    // Output = state[0..4] || state[12..16] (no initial-state add).
    let mut out = [0u8; 32];
    for i in 0..4 {
        out[4 * i..4 * i + 4].copy_from_slice(&state[i].to_le_bytes());
    }
    for i in 0..4 {
        out[16 + 4 * i..16 + 4 * i + 4].copy_from_slice(&state[12 + i].to_le_bytes());
    }
    out
}

// ============================================================================
// XChaCha20-Poly1305 AEAD
// ============================================================================

/// XChaCha20-Poly1305 AEAD with 24-byte nonce.
///
/// Same shape as [`ChaCha20Poly1305`]:
/// a unit struct with associated `encrypt` / `decrypt` functions.
/// The only difference is the 24-byte nonce.
pub struct XChaCha20Poly1305;

impl XChaCha20Poly1305 {
    /// Encrypt and authenticate a message.
    ///
    /// Returns `(ciphertext, tag)` where `ciphertext.len() ==
    /// plaintext.len()` and `tag` is exactly 16 bytes.
    pub fn encrypt(key: &[u8; 32], nonce: &[u8; 24], aad: &[u8], plaintext: &[u8]) -> (Vec<u8>, [u8; 16]) {
        let (subkey, nonce12) = derive(key, nonce);
        ChaCha20Poly1305::encrypt(&subkey, &nonce12, aad, plaintext)
    }

    /// Decrypt and verify a ciphertext. Returns `None` on tag mismatch.
    pub fn decrypt(key: &[u8; 32], nonce: &[u8; 24], aad: &[u8], ciphertext: &[u8], tag: &[u8; 16]) -> Option<Vec<u8>> {
        let (subkey, nonce12) = derive(key, nonce);
        ChaCha20Poly1305::decrypt(&subkey, &nonce12, aad, ciphertext, tag)
    }
}

/// Derive the inner (subkey, 12-byte nonce) pair from the 24-byte nonce.
fn derive(key: &[u8; 32], nonce: &[u8; 24]) -> ([u8; 32], [u8; 12]) {
    let mut hchacha_in = [0u8; 16];
    hchacha_in.copy_from_slice(&nonce[..16]);
    let subkey = hchacha20(key, &hchacha_in);

    let mut nonce12 = [0u8; 12];
    // Prefix of 4 zero bytes + last 8 bytes of the 24-byte nonce.
    nonce12[4..].copy_from_slice(&nonce[16..24]);

    (subkey, nonce12)
}

// ============================================================================
// Tests
// ============================================================================

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

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

    /// HChaCha20 test vector from draft-irtf-cfrg-xchacha §2.2.1.
    #[test]
    fn hchacha20_test_vector() {
        let key = hex("00 01 02 03 04 05 06 07 08 09 0a 0b 0c 0d 0e 0f
             10 11 12 13 14 15 16 17 18 19 1a 1b 1c 1d 1e 1f");
        let input = hex("00 00 00 09 00 00 00 4a 00 00 00 00 31 41 59 27");
        let expected = hex("82 41 3b 42 27 b2 7b fe d3 0e 42 50 8a 87 7d 73
             a0 f9 e4 d5 8a 74 a8 53 c1 2e c4 13 26 d3 ec dc");
        let k: [u8; 32] = key.try_into().unwrap();
        let n: [u8; 16] = input.try_into().unwrap();
        let out = hchacha20(&k, &n);
        assert_eq!(out.to_vec(), expected);
    }

    /// XChaCha20-Poly1305 test vector from draft-irtf-cfrg-xchacha §A.3.1.
    #[test]
    fn xchacha20poly1305_test_vector() {
        let plaintext = hex("4c 61 64 69 65 73 20 61 6e 64 20 47 65 6e 74 6c
             65 6d 65 6e 20 6f 66 20 74 68 65 20 63 6c 61 73
             73 20 6f 66 20 27 39 39 3a 20 49 66 20 49 20 63
             6f 75 6c 64 20 6f 66 66 65 72 20 79 6f 75 20 6f
             6e 6c 79 20 6f 6e 65 20 74 69 70 20 66 6f 72 20
             74 68 65 20 66 75 74 75 72 65 2c 20 73 75 6e 73
             63 72 65 65 6e 20 77 6f 75 6c 64 20 62 65 20 69
             74 2e");
        let aad = hex("50 51 52 53 c0 c1 c2 c3 c4 c5 c6 c7");
        let key = hex("80 81 82 83 84 85 86 87 88 89 8a 8b 8c 8d 8e 8f
             90 91 92 93 94 95 96 97 98 99 9a 9b 9c 9d 9e 9f");
        let nonce = hex("40 41 42 43 44 45 46 47 48 49 4a 4b 4c 4d 4e 4f
             50 51 52 53 54 55 56 57");
        let expected_ct = hex("bd 6d 17 9d 3e 83 d4 3b 95 76 57 94 93 c0 e9 39
             57 2a 17 00 25 2b fa cc be d2 90 2c 21 39 6c bb
             73 1c 7f 1b 0b 4a a6 44 0b f3 a8 2f 4e da 7e 39
             ae 64 c6 70 8c 54 c2 16 cb 96 b7 2e 12 13 b4 52
             2f 8c 9b a4 0d b5 d9 45 b1 1b 69 b9 82 c1 bb 9e
             3f 3f ac 2b c3 69 48 8f 76 b2 38 35 65 d3 ff f9
             21 f9 66 4c 97 63 7d a9 76 88 12 f6 15 c6 8b 13
             b5 2e");
        let expected_tag = hex("c0 87 59 24 c1 c7 98 79 47 de af d8 78 0a cf 49");

        let k: [u8; 32] = key.try_into().unwrap();
        let n: [u8; 24] = nonce.try_into().unwrap();

        let (ct, tag) = XChaCha20Poly1305::encrypt(&k, &n, &aad, &plaintext);
        assert_eq!(ct, expected_ct, "ciphertext mismatch");
        assert_eq!(tag.to_vec(), expected_tag, "tag mismatch");

        let pt = XChaCha20Poly1305::decrypt(&k, &n, &aad, &ct, &tag).expect("decrypt must succeed");
        assert_eq!(pt, plaintext);
    }

    #[test]
    fn xchacha20poly1305_tamper_rejected() {
        let key = [0x42u8; 32];
        let nonce = [0x77u8; 24];
        let aad = b"header";
        let pt = b"secret payload";

        let (mut ct, tag) = XChaCha20Poly1305::encrypt(&key, &nonce, aad, pt);
        // Flip one ciphertext byte → decrypt must reject.
        ct[0] ^= 0xFF;
        assert!(XChaCha20Poly1305::decrypt(&key, &nonce, aad, &ct, &tag).is_none());
    }

    #[test]
    fn xchacha20poly1305_empty_plaintext() {
        let key = [0u8; 32];
        let nonce = [0u8; 24];
        let (ct, tag) = XChaCha20Poly1305::encrypt(&key, &nonce, b"", b"");
        assert_eq!(ct.len(), 0);
        let pt = XChaCha20Poly1305::decrypt(&key, &nonce, b"", &ct, &tag).unwrap();
        assert_eq!(pt.len(), 0);
    }
}
