//! Reticulum's modified Fernet token (`ARCHITECTURE.md` §3.3).

use aes::Aes256;
use cbc::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use hmac::{Hmac, Mac};
use rand_core::CryptoRngCore;
use sha2::Sha256;
use zeroize::ZeroizeOnDrop;

use crate::derive::DerivedKey;
use crate::error::Error;

type Aes256CbcEnc = cbc::Encryptor<Aes256>;
type Aes256CbcDec = cbc::Decryptor<Aes256>;
type HmacSha256 = Hmac<Sha256>;

const IV_LEN: usize = 16;
const MAC_LEN: usize = 32;
const BLOCK_LEN: usize = 16;

/// Reticulum's modified Fernet:
/// `IV(16) || AES-256-CBC(PKCS7(plaintext)) || HMAC-SHA-256(IV || CT)`.
///
/// There is **no** version byte and **no** timestamp — deliberately, to avoid
/// metadata leakage. HMAC verification is constant-time
/// (`subtle::ConstantTimeEq`, via `hmac`'s `verify_slice`); the Python
/// reference's `==` on `bytes` is not, and we fix that without breaking wire
/// compatibility (§3.3, §9.2).
#[derive(ZeroizeOnDrop)]
pub struct Token {
    encryption_key: [u8; 32],
    signing_key: [u8; 32],
}

impl Token {
    /// Per-message overhead added to every ciphertext: `IV(16) + HMAC(32)`.
    pub const OVERHEAD: usize = IV_LEN + MAC_LEN;

    /// Construct a token from a 64-byte derived key (§3.3 split convention).
    #[must_use]
    pub fn new(key: &DerivedKey) -> Self {
        Self {
            signing_key: *key.signing_half(),
            encryption_key: *key.encryption_half(),
        }
    }

    /// Encrypt `plaintext` into `out`, returning the number of bytes written.
    ///
    /// # Errors
    /// Returns [`Error::InvalidLength`] if `out` is too small for
    /// `OVERHEAD + padded(plaintext)`, or [`Error::InvalidKey`] on key setup
    /// failure.
    pub fn encrypt<R: CryptoRngCore>(
        &self,
        rng: &mut R,
        plaintext: &[u8],
        out: &mut [u8],
    ) -> Result<usize, Error> {
        let mut iv = [0u8; IV_LEN];
        rng.fill_bytes(&mut iv);

        let ct = Aes256CbcEnc::new_from_slices(&self.encryption_key, &iv)
            .map_err(|_| Error::InvalidKey)?
            .encrypt_padded_vec_mut::<Pkcs7>(plaintext);

        let total = IV_LEN + ct.len() + MAC_LEN;
        if out.len() < total {
            return Err(Error::InvalidLength);
        }

        let mut mac =
            HmacSha256::new_from_slice(&self.signing_key).map_err(|_| Error::InvalidKey)?;
        mac.update(&iv);
        mac.update(&ct);
        let tag = mac.finalize().into_bytes();

        out[..IV_LEN].copy_from_slice(&iv);
        out[IV_LEN..IV_LEN + ct.len()].copy_from_slice(&ct);
        out[IV_LEN + ct.len()..total].copy_from_slice(tag.as_slice());
        Ok(total)
    }

    /// Decrypt `ciphertext` into `out`, returning the plaintext slice.
    ///
    /// Verifies the HMAC in constant time **before** decrypting; a bad MAC
    /// yields [`Error::MacMismatch`] and no plaintext.
    ///
    /// # Errors
    /// [`Error::InvalidLength`] for a malformed length or too-small `out`,
    /// [`Error::MacMismatch`] for a failed MAC, [`Error::InvalidKey`] on key
    /// setup failure.
    pub fn decrypt<'a>(&self, ciphertext: &[u8], out: &'a mut [u8]) -> Result<&'a [u8], Error> {
        if ciphertext.len() < Self::OVERHEAD + BLOCK_LEN
            || (ciphertext.len() - Self::OVERHEAD) % BLOCK_LEN != 0
        {
            return Err(Error::InvalidLength);
        }
        let (iv, rest) = ciphertext.split_at(IV_LEN);
        let (ct, mac) = rest.split_at(rest.len() - MAC_LEN);

        let mut h = HmacSha256::new_from_slice(&self.signing_key).map_err(|_| Error::InvalidKey)?;
        h.update(iv);
        h.update(ct);
        h.verify_slice(mac).map_err(|_| Error::MacMismatch)?;

        let pt = Aes256CbcDec::new_from_slices(&self.encryption_key, iv)
            .map_err(|_| Error::InvalidKey)?
            .decrypt_padded_vec_mut::<Pkcs7>(ct)
            .map_err(|_| Error::InvalidLength)?;

        if out.len() < pt.len() {
            return Err(Error::InvalidLength);
        }
        out[..pt.len()].copy_from_slice(&pt);
        Ok(&out[..pt.len()])
    }
}

#[cfg(test)]
mod tests {
    use super::Token;
    use crate::derive::DerivedKey;
    use crate::error::Error;
    use rand_core::OsRng;

    fn token() -> Token {
        // Distinct halves so a swapped split would be caught.
        let mut key = [0u8; 64];
        key[..32].fill(0xAA);
        key[32..].fill(0xBB);
        Token::new(&DerivedKey::from_bytes(key))
    }

    #[test]
    fn overhead_is_iv_plus_mac() {
        assert_eq!(Token::OVERHEAD, 48);
    }

    #[test]
    fn round_trip() {
        let t = token();
        let pt = b"hello reticulum, this spans multiple AES blocks!";
        let mut ct = [0u8; 256];
        let mut rng = OsRng;
        let n = t.encrypt(&mut rng, pt, &mut ct).unwrap();
        let mut out = [0u8; 256];
        assert_eq!(t.decrypt(&ct[..n], &mut out).unwrap(), pt);
    }

    #[test]
    fn round_trip_empty_plaintext() {
        let t = token();
        let mut ct = [0u8; 64];
        let mut rng = OsRng;
        let n = t.encrypt(&mut rng, b"", &mut ct).unwrap();
        // empty plaintext still pads to one block under PKCS7.
        assert_eq!(n, Token::OVERHEAD + 16);
        let mut out = [0u8; 64];
        assert_eq!(t.decrypt(&ct[..n], &mut out).unwrap(), b"");
    }

    #[test]
    fn tampered_mac_rejected() {
        let t = token();
        let mut ct = [0u8; 128];
        let mut rng = OsRng;
        let n = t.encrypt(&mut rng, b"data", &mut ct).unwrap();
        ct[n - 1] ^= 0x01; // flip a MAC byte
        let mut out = [0u8; 128];
        assert_eq!(t.decrypt(&ct[..n], &mut out), Err(Error::MacMismatch));
    }

    #[test]
    fn tampered_ciphertext_rejected() {
        let t = token();
        let mut ct = [0u8; 128];
        let mut rng = OsRng;
        let n = t.encrypt(&mut rng, b"data", &mut ct).unwrap();
        ct[16] ^= 0x01; // flip the first ciphertext byte (after the 16-byte IV) → MAC fails
        let mut out = [0u8; 128];
        assert_eq!(t.decrypt(&ct[..n], &mut out), Err(Error::MacMismatch));
    }

    #[test]
    fn short_ciphertext_rejected() {
        let t = token();
        let mut out = [0u8; 64];
        assert_eq!(t.decrypt(&[0u8; 10], &mut out), Err(Error::InvalidLength));
    }
}
