//! Authenticated encryption session built on ChaCha20-Poly1305.
//!
//! An [`AeadSession`] wraps a 256-bit key and a 64-bit message counter. Each [`AeadSession::seal`]
//! / [`AeadSession::open`] call derives a unique 96-bit nonce from the counter and then advances
//! it, so a single session never reuses a nonce (until the 2^64 counter wraps, which is
//! unreachable in practice).
//!
//! The two endpoints of a connection keep one session per direction (see
//! [`crate::SessionKeys`]). A sender's `seal` session and the matching receiver's `open` session
//! advance their counters in lockstep, so they stay aligned without transmitting the nonce.

use chacha20poly1305::aead::{Aead, KeyInit, Payload};
use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
use zeroize::Zeroize;

use crate::CryptoError;

/// A directional AEAD session: a key plus a monotonically increasing nonce counter.
pub struct AeadSession {
    key: [u8; 32],
    counter: u64,
}

impl AeadSession {
    /// Create a new session from a 256-bit key, starting at counter 0.
    #[must_use]
    pub fn new(key: [u8; 32]) -> Self {
        Self { key, counter: 0 }
    }

    /// Derive the 96-bit (12-byte) nonce for a given counter value.
    ///
    /// Layout: little-endian `u64` counter in bytes `[0..8]`, then four zero bytes in `[8..12]`.
    /// Exposed (crate-internal) so tests can assert nonce uniqueness directly.
    #[must_use]
    pub(crate) fn nonce_for(counter: u64) -> [u8; 12] {
        let mut nonce = [0u8; 12];
        nonce[..8].copy_from_slice(&counter.to_le_bytes());
        // bytes [8..12] stay zero
        nonce
    }

    /// Build the cipher instance for the current key.
    fn cipher(&self) -> ChaCha20Poly1305 {
        ChaCha20Poly1305::new(Key::from_slice(&self.key))
    }

    /// Encrypt `plaintext` with associated data `aad`, returning `ciphertext || tag`.
    ///
    /// Uses the current counter as the nonce and then increments it.
    ///
    /// # Panics
    /// Panics if the 64-bit message counter overflows (after 2^64 messages on one key) or if the
    /// underlying AEAD reports an error (which, for ChaCha20-Poly1305 encryption, only happens
    /// when the plaintext exceeds the cipher's maximum supported length).
    pub fn seal(&mut self, plaintext: &[u8], aad: &[u8]) -> Vec<u8> {
        let nonce = Self::nonce_for(self.counter);
        let ct = self
            .cipher()
            .encrypt(
                Nonce::from_slice(&nonce),
                Payload {
                    msg: plaintext,
                    aad,
                },
            )
            .expect("ChaCha20-Poly1305 encryption never fails for in-range plaintext");
        self.counter = self
            .counter
            .checked_add(1)
            .expect("AEAD nonce counter overflow");
        ct
    }

    /// Decrypt `ciphertext` (which must be `ciphertext || tag`) with associated data `aad`.
    ///
    /// Uses the current counter as the nonce and then increments it (symmetrically to
    /// [`AeadSession::seal`]), so a paired seal/open pair of sessions stay aligned even across a
    /// failed decryption.
    ///
    /// # Errors
    /// Returns [`CryptoError::AeadDecrypt`] if authentication fails (tampered ciphertext, wrong
    /// AAD, wrong key, or desynchronized counter).
    pub fn open(&mut self, ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>, CryptoError> {
        let nonce = Self::nonce_for(self.counter);
        let result = self.cipher().decrypt(
            Nonce::from_slice(&nonce),
            Payload {
                msg: ciphertext,
                aad,
            },
        );
        // Advance symmetrically to `seal`, regardless of success, to keep counters aligned.
        self.counter = self
            .counter
            .checked_add(1)
            .expect("AEAD nonce counter overflow");
        result.map_err(|_| CryptoError::AeadDecrypt)
    }

    /// Consume the session, returning a reusable explicit-nonce [`AeadKey`] plus the current
    /// counter value. Datagram transports use this to continue from the post-handshake counter
    /// while carrying the nonce on the wire (see [`AeadKey`]).
    #[must_use]
    pub fn into_parts(self) -> (AeadKey, u64) {
        (AeadKey::new(self.key), self.counter)
    }

    /// Current counter value (next nonce to be used). Test-only accessor.
    #[cfg(test)]
    #[must_use]
    pub(crate) fn counter(&self) -> u64 {
        self.counter
    }
}

impl Drop for AeadSession {
    fn drop(&mut self) {
        self.key.zeroize();
    }
}

impl zeroize::ZeroizeOnDrop for AeadSession {}

/// A 256-bit ChaCha20-Poly1305 key used with EXPLICIT per-message nonces.
///
/// Unlike [`AeadSession`] (which derives the nonce from an internal, lock-step counter), `AeadKey`
/// takes the nonce counter as an argument on every call. Datagram transports need exactly this:
/// packets may be lost or reordered, so the per-record counter is carried on the wire and supplied
/// by the caller rather than tracked implicitly. The nonce scheme is identical to [`AeadSession`]
/// (`LE(counter) || 0x0000_0000`), so the two interoperate on the same key as long as their
/// counter ranges do not overlap.
pub struct AeadKey {
    key: [u8; 32],
}

impl AeadKey {
    /// Create a key holder from 256 bits of key material.
    #[must_use]
    pub fn new(key: [u8; 32]) -> Self {
        Self { key }
    }

    /// Build the cipher instance for the current key.
    fn cipher(&self) -> ChaCha20Poly1305 {
        ChaCha20Poly1305::new(Key::from_slice(&self.key))
    }

    /// Encrypt `plaintext` with associated data `aad` under the nonce derived from `counter`,
    /// returning `ciphertext || tag`. The caller owns nonce uniqueness: never reuse a `counter`
    /// value with the same key.
    ///
    /// # Panics
    /// Panics only if the underlying AEAD reports an error, which for ChaCha20-Poly1305 encryption
    /// happens solely when the plaintext exceeds the cipher's maximum supported length.
    #[must_use]
    pub fn seal(&self, counter: u64, plaintext: &[u8], aad: &[u8]) -> Vec<u8> {
        let nonce = AeadSession::nonce_for(counter);
        self.cipher()
            .encrypt(
                Nonce::from_slice(&nonce),
                Payload {
                    msg: plaintext,
                    aad,
                },
            )
            .expect("ChaCha20-Poly1305 encryption never fails for in-range plaintext")
    }

    /// Decrypt `ciphertext` (`ciphertext || tag`) with associated data `aad` under the nonce
    /// derived from `counter`.
    ///
    /// # Errors
    /// Returns [`CryptoError::AeadDecrypt`] if authentication fails (tampered ciphertext, wrong
    /// AAD, wrong key, or wrong counter).
    pub fn open(
        &self,
        counter: u64,
        ciphertext: &[u8],
        aad: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        let nonce = AeadSession::nonce_for(counter);
        self.cipher()
            .decrypt(
                Nonce::from_slice(&nonce),
                Payload {
                    msg: ciphertext,
                    aad,
                },
            )
            .map_err(|_| CryptoError::AeadDecrypt)
    }
}

impl Drop for AeadKey {
    fn drop(&mut self) {
        self.key.zeroize();
    }
}

impl zeroize::ZeroizeOnDrop for AeadKey {}

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

    #[test]
    fn nonce_layout_is_le_counter_then_zeros() {
        assert_eq!(AeadSession::nonce_for(0), [0u8; 12]);

        let mut expected = [0u8; 12];
        expected[0] = 1;
        assert_eq!(AeadSession::nonce_for(1), expected);

        // 0x0102030405060708 little-endian in the first 8 bytes, zeros after.
        let n = AeadSession::nonce_for(0x0807_0605_0403_0201);
        assert_eq!(&n[..8], &[1, 2, 3, 4, 5, 6, 7, 8]);
        assert_eq!(&n[8..], &[0, 0, 0, 0]);
    }

    #[test]
    fn counter_is_monotonic_per_seal() {
        let mut s = AeadSession::new([0u8; 32]);
        assert_eq!(s.counter(), 0);
        for expected_next in 1..=64u64 {
            let _ = s.seal(b"x", b"");
            assert_eq!(s.counter(), expected_next);
        }
    }

    #[test]
    fn nonces_are_distinct_over_10_000_counters() {
        // Directly exercise the nonce derivation (the crate-internal, testable surface).
        let mut seen: HashSet<[u8; 12]> = HashSet::with_capacity(10_000);
        for c in 0..10_000u64 {
            assert!(
                seen.insert(AeadSession::nonce_for(c)),
                "duplicate nonce at {c}"
            );
        }
        assert_eq!(seen.len(), 10_000);
    }

    #[test]
    fn aead_key_explicit_nonce_roundtrip() {
        let k = AeadKey::new([7u8; 32]);
        let ct = k.seal(42, b"hello datagram", b"aad");
        let pt = k.open(42, &ct, b"aad").expect("open at same counter");
        assert_eq!(pt, b"hello datagram");
    }

    #[test]
    fn aead_key_wrong_counter_or_aad_fails() {
        let k = AeadKey::new([9u8; 32]);
        let ct = k.seal(5, b"msg", b"aad");
        assert!(k.open(6, &ct, b"aad").is_err(), "wrong counter must fail");
        assert!(k.open(5, &ct, b"other").is_err(), "wrong aad must fail");
        let mut tampered = ct.clone();
        tampered[0] ^= 1;
        assert!(k.open(5, &tampered, b"aad").is_err(), "tamper must fail");
    }

    #[test]
    fn aead_key_matches_session_nonce_scheme() {
        // An AeadKey at counter c must reproduce what an AeadSession produces at counter c.
        let key = [3u8; 32];
        let mut s = AeadSession::new(key);
        let from_session = s.seal(b"x", b"a"); // counter 0, then advances
        let from_key = AeadKey::new(key).seal(0, b"x", b"a");
        assert_eq!(from_session, from_key);
    }

    #[test]
    fn into_parts_preserves_key_and_counter() {
        let mut s = AeadSession::new([5u8; 32]);
        let _ = s.seal(b"a", b"");
        let _ = s.seal(b"b", b""); // counter now 2
        let (key, counter) = s.into_parts();
        assert_eq!(counter, 2);
        // The recovered key, used at the next counter, matches a fresh session advanced to 2.
        let expect = {
            let mut s2 = AeadSession::new([5u8; 32]);
            let _ = s2.seal(b"a", b"");
            let _ = s2.seal(b"b", b"");
            s2.seal(b"c", b"d")
        };
        assert_eq!(key.seal(counter, b"c", b"d"), expect);
    }
}