//! ChaCha20 stream cipher (RFC 8439).
//!
//! This is the IETF / TLS 1.3 variant of ChaCha20: 256-bit key,
//! 96-bit nonce, 32-bit block counter, 64-byte block size, 20 rounds.
//!
//! It is the second AEAD primitive shipped by `arcana`
//! alongside AES-GCM. Used by TLS 1.3, Noise, Signal, WireGuard,
//! QUIC, OpenSSH, and most modern protocols that prefer a constant-
//! time stream cipher with no S-box dependencies (no cache-timing
//! surface, in contrast to table-based AES).
//!
//! # Layout
//!
//! ```text
//! state (4x4 u32 little-endian):
//!
//!   constants  constants  constants  constants    "expand 32-byte k"
//!   key        key        key        key
//!   key        key        key        key
//!   counter    nonce      nonce      nonce
//! ```
//!
//! Each 64-byte block is computed as `serialize(rounds(state) + state)`.
//! Successive blocks increment `counter`.
//!
//! # API
//!
//! ```rust,ignore
//! use arcana::cipher::chacha20::ChaCha20;
//!
//! let mut cipher = ChaCha20::new(&key, &nonce, 1); // initial counter = 1
//! let mut buf = b"plaintext".to_vec();
//! cipher.apply_keystream(&mut buf);  // encrypt or decrypt
//! ```
//!
//! Stream ciphers are symmetric: `apply_keystream` does both
//! encryption and decryption since the keystream is XOR'd with
//! whatever is passed in.
//!
//! # Tests
//!
//! Pinned against RFC 8439 §2.3.2 (block test vector) and §2.4.2
//! (encryption test vector).

// ============================================================================
// Quarter-round (RFC 8439 §2.1)
// ============================================================================

/// Single ChaCha20 quarter-round on 4 state words. Operates in place.
///
/// ```text
///   a += b; d ^= a; d <<<= 16
///   c += d; b ^= c; b <<<= 12
///   a += b; d ^= a; d <<<=  8
///   c += d; b ^= c; b <<<=  7
/// ```
#[inline(always)]
pub(crate) 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);
}

// ============================================================================
// Block function (RFC 8439 §2.3)
// ============================================================================

/// Compute one 64-byte ChaCha20 keystream block from the initial state.
///
/// 20 rounds = 10 double-rounds; each double-round is 4 column rounds
/// followed by 4 diagonal rounds. The output is `(working ⊞ initial)`
/// serialized little-endian, where `⊞` is wrapping word-wise add.
fn block(state: &[u32; 16]) -> [u8; 64] {
    let mut working = *state;

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

    // Output = working + state, serialized little-endian.
    let mut out = [0u8; 64];
    for i in 0..16 {
        let word = working[i].wrapping_add(state[i]);
        out[4 * i..4 * i + 4].copy_from_slice(&word.to_le_bytes());
    }
    out
}

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

/// ChaCha20 stream cipher state (RFC 8439).
///
/// Holds the initial state (key + nonce + counter) and the current
/// counter value. Buffers the leftover bytes from the previous
/// keystream block so partial calls to `apply_keystream` work
/// correctly.
#[derive(Clone)]
pub struct ChaCha20 {
    /// Initial 16-word state. The counter at index 12 is mutated
    /// in place between blocks.
    state: [u32; 16],
    /// Current 64-byte keystream block, possibly partially consumed.
    buffer: [u8; 64],
    /// Number of bytes left in `buffer` (0..=64).
    buf_pos: usize,
}

impl ChaCha20 {
    /// Initialise a ChaCha20 cipher with a 32-byte key, a 12-byte
    /// nonce, and an initial 32-bit block counter.
    ///
    /// Per RFC 8439 §2.4 the AEAD construction starts the cipher
    /// at counter = 1 (counter = 0 produces the one-time Poly1305
    /// key). Standalone ChaCha20 users typically start at counter
    /// = 0 or 1 -- whatever their protocol specifies.
    pub fn new(key: &[u8; 32], nonce: &[u8; 12], counter: u32) -> Self {
        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());
        }
        // Counter (1 word).
        state[12] = counter;
        // Nonce (3 words, LE).
        for i in 0..3 {
            state[13 + i] = u32::from_le_bytes(nonce[4 * i..4 * i + 4].try_into().unwrap());
        }

        Self {
            state,
            buffer: [0u8; 64],
            buf_pos: 64,
        } // buf_pos = 64 means "empty"
    }

    /// XOR the keystream into `data` in place. Encrypts or decrypts
    /// indifferently (the cipher is symmetric).
    ///
    /// Handles arbitrary lengths and partial blocks: the cipher
    /// remembers leftover keystream bytes between calls, so
    /// `cipher.apply_keystream(b"hi"); cipher.apply_keystream(b"!")`
    /// produces the same output as one call with `b"hi!"`.
    pub fn apply_keystream(&mut self, data: &mut [u8]) {
        let mut pos = 0;
        while pos < data.len() {
            // Refill the buffer if it is empty.
            if self.buf_pos == 64 {
                self.buffer = block(&self.state);
                // Advance the 32-bit block counter (RFC 8439 §2.3).
                self.state[12] = self.state[12].wrapping_add(1);
                self.buf_pos = 0;
            }
            let take = (64 - self.buf_pos).min(data.len() - pos);
            for i in 0..take {
                data[pos + i] ^= self.buffer[self.buf_pos + i];
            }
            self.buf_pos += take;
            pos += take;
        }
    }
}

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

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

    fn hex(s: &str) -> Vec<u8> {
        assert!(s.len() % 2 == 0);
        (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.1.1 quarter-round example.
    ///
    /// Input:  a=0x11111111, b=0x01020304, c=0x9b8d6f43, d=0x01234567
    /// Output: a=0xea2a92f4, b=0xcb1cf8ce, c=0x4581472e, d=0x5881c4bb
    #[test]
    fn rfc8439_2_1_1_quarter_round() {
        let mut s = [0u32; 16];
        s[0] = 0x11111111;
        s[1] = 0x01020304;
        s[2] = 0x9b8d6f43;
        s[3] = 0x01234567;
        quarter_round(&mut s, 0, 1, 2, 3);
        assert_eq!(s[0], 0xea2a92f4);
        assert_eq!(s[1], 0xcb1cf8ce);
        assert_eq!(s[2], 0x4581472e);
        assert_eq!(s[3], 0x5881c4bb);
    }

    /// RFC 8439 §2.3.2 block-function test vector.
    ///
    /// Key:        00:01:02:..:1f
    /// Nonce:      00:00:00:09:00:00:00:4a:00:00:00:00
    /// Counter:    1
    /// Output:     10f1e7e4d13b5915500fdd1fa32071c4
    ///             c7d1f4c733c068030422aa9ac3d46c4e
    ///             d2826446079faa0914c2d705d98b02a2
    ///             b5129cd1de164eb9cbd083e8a2503c4e
    #[test]
    fn rfc8439_2_3_2_block_vector() {
        let key: [u8; 32] = hex_arr("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f");
        let nonce: [u8; 12] = hex_arr("000000090000004a00000000");
        let cipher = ChaCha20::new(&key, &nonce, 1);
        let out = block(&cipher.state);
        let expected = hex("10f1e7e4d13b5915500fdd1fa32071c4\
             c7d1f4c733c068030422aa9ac3d46c4e\
             d2826446079faa0914c2d705d98b02a2\
             b5129cd1de164eb9cbd083e8a2503c4e");
        assert_eq!(out.to_vec(), expected);
    }

    /// RFC 8439 §2.4.2 encryption test vector.
    ///
    /// Key:        00:01:02:..:1f
    /// Nonce:      00:00:00:00:00:00:00:4a:00:00:00:00
    /// Counter:    1
    /// Plaintext:  "Ladies and Gentlemen of the class of '99: ...
    ///              ...if I could offer you only one tip for the
    ///              future, sunscreen would be it."
    /// Ciphertext: 6e2e359a2568f98041ba0728dd0d6981...
    ///             ...874d
    #[test]
    fn rfc8439_2_4_2_encryption_vector() {
        let key: [u8; 32] = hex_arr("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f");
        let nonce: [u8; 12] = hex_arr("000000000000004a00000000");
        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 mut buf = plaintext.to_vec();
        let mut cipher = ChaCha20::new(&key, &nonce, 1);
        cipher.apply_keystream(&mut buf);

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

    /// Decryption is the same operation as encryption (XOR keystream).
    /// Re-applying the keystream must recover the plaintext.
    #[test]
    fn chacha20_encrypt_decrypt_roundtrip() {
        let key = [0x42u8; 32];
        let nonce = [0xa5u8; 12];
        let plaintext = b"the quick brown fox jumps over the lazy dog";
        let mut buf = plaintext.to_vec();

        let mut enc = ChaCha20::new(&key, &nonce, 1);
        enc.apply_keystream(&mut buf);
        assert_ne!(buf.as_slice(), plaintext);

        let mut dec = ChaCha20::new(&key, &nonce, 1);
        dec.apply_keystream(&mut buf);
        assert_eq!(buf.as_slice(), plaintext);
    }

    /// Calling `apply_keystream` in chunks must give the same result
    /// as one big call (proves the buffered partial-block path is
    /// stateful and correct).
    #[test]
    fn chacha20_chunked_apply_matches_single() {
        let key = [0x77u8; 32];
        let nonce = [0x11u8; 12];
        let plaintext: Vec<u8> = (0..200).map(|i| i as u8).collect();

        // Single call.
        let mut single = plaintext.clone();
        ChaCha20::new(&key, &nonce, 1).apply_keystream(&mut single);

        // Chunked at boundaries that span block edges (64).
        let mut chunked = plaintext.clone();
        let mut c = ChaCha20::new(&key, &nonce, 1);
        c.apply_keystream(&mut chunked[..30]);
        c.apply_keystream(&mut chunked[30..70]); // crosses block boundary
        c.apply_keystream(&mut chunked[70..130]); // crosses again
        c.apply_keystream(&mut chunked[130..]);

        assert_eq!(chunked, single);
    }

    /// All-zero key + all-zero nonce + counter 0 yields a known
    /// fixed first block (RFC 8439 §2.3.2 sanity).
    #[test]
    fn chacha20_zero_key_zero_nonce_block_zero() {
        let key = [0u8; 32];
        let nonce = [0u8; 12];
        let cipher = ChaCha20::new(&key, &nonce, 0);
        let out = block(&cipher.state);
        // First 32 bytes from RFC 8439 Appendix A.1 test vector 1.
        let expected_prefix = hex("76b8e0ada0f13d90405d6ae55386bd28\
             bdd219b8a08ded1aa836efcc8b770dc7");
        assert_eq!(&out[..32], expected_prefix.as_slice());
    }
}
