#![doc = include_str!("../README.md")]

use aes_gcm::aead::{Aead, Generate};
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use anyhow::{Context, Result};
use secrecy::SecretString;

/// A struct that encapsulates token encryption and decryption using AES-256-GCM.
pub struct TokenEncryption {
    cipher: Aes256Gcm,
}

impl TokenEncryption {
    /// Creates a new [`TokenEncryption`] instance with the provided cipher
    pub fn new(cipher: Aes256Gcm) -> Self {
        Self { cipher }
    }

    /// Creates a new [`TokenEncryption`] instance with a cipher for testing
    /// purposes.
    #[cfg(any(test, debug_assertions))]
    pub fn for_testing() -> Self {
        let test_key = b"test_key_32_bytes_long_for_tests";
        Self::new(Aes256Gcm::new_from_slice(test_key).unwrap())
    }

    /// Creates a new [`TokenEncryption`] instance from the environment
    ///
    /// Reads the `TOKEN_ENCRYPTION_KEY` environment variable, which
    /// should be a 64-character hex string (32 bytes when decoded).
    pub fn from_environment() -> Result<Self> {
        let token_key = std::env::var("TOKEN_ENCRYPTION_KEY")
            .context("TOKEN_ENCRYPTION_KEY environment variable not set")?;

        if token_key.len() != 64 {
            anyhow::bail!("TOKEN_ENCRYPTION_KEY must be exactly 64 hex characters");
        }

        let token_key = hex::decode(token_key.as_bytes())
            .context("TOKEN_ENCRYPTION_KEY must be exactly 64 hex characters")?;

        let cipher = Aes256Gcm::new_from_slice(&token_key)
            .context("TOKEN_ENCRYPTION_KEY must be exactly 64 hex characters")?;

        Ok(Self::new(cipher))
    }

    /// Encrypts a token using AES-256-GCM
    ///
    /// The encrypted data format is: `[12-byte nonce][encrypted data]`
    /// The nonce is randomly generated for each encryption to ensure uniqueness.
    pub fn encrypt(&self, plaintext: &str) -> Result<Vec<u8>> {
        // Generate a random nonce for this encryption
        let nonce = Nonce::generate();

        // Encrypt the token
        let encrypted = self
            .cipher
            .encrypt(&nonce, plaintext.as_bytes())
            .map_err(|error| anyhow::anyhow!("Failed to encrypt token: {error}"))?;

        // Combine nonce + ciphertext (which includes the auth tag)
        let mut result = Vec::with_capacity(nonce.len() + encrypted.len());
        result.extend_from_slice(&nonce);
        result.extend_from_slice(&encrypted);

        Ok(result)
    }

    /// Decrypts a token using AES-256-GCM
    ///
    /// Expects the data format: `[12-byte nonce][encrypted data]`
    pub fn decrypt(&self, encrypted: &[u8]) -> Result<SecretString> {
        if encrypted.len() < 12 {
            anyhow::bail!("Invalid encrypted token: too short");
        }

        // Extract nonce and ciphertext
        let (nonce_bytes, ciphertext) = encrypted.split_at(12);
        let nonce = Nonce::try_from(nonce_bytes).context("Invalid encrypted token: bad nonce")?;

        // Decrypt the token
        let plaintext = self
            .cipher
            .decrypt(&nonce, ciphertext)
            .context("Failed to decrypt token")?;

        let plaintext =
            String::from_utf8(plaintext).context("Decrypted token is not valid UTF-8")?;

        Ok(plaintext.into())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aes_gcm::KeyInit;
    use claims::{assert_err, assert_ok};
    use insta::assert_snapshot;
    use secrecy::ExposeSecret;

    fn create_test_encryption() -> TokenEncryption {
        let cipher = Aes256Gcm::new_from_slice(b"test_master_key_32_bytes_long!!!").unwrap();
        TokenEncryption { cipher }
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let encryption = create_test_encryption();
        let original_token = "ghs_test_token_123456789";

        // Encrypt the token
        let encrypted = assert_ok!(encryption.encrypt(original_token));

        // Decrypt it back
        let decrypted = assert_ok!(encryption.decrypt(&encrypted));

        assert_eq!(original_token, decrypted.expose_secret());
    }

    #[test]
    fn test_encrypt_produces_different_ciphertext() {
        let encryption = create_test_encryption();
        let token = "ghs_test_token_123456789";

        // Encrypt the same token twice
        let encrypted1 = assert_ok!(encryption.encrypt(token));
        let encrypted2 = assert_ok!(encryption.encrypt(token));

        // Should produce different ciphertext due to random nonce
        assert_ne!(encrypted1, encrypted2);

        // But both should decrypt to the same plaintext
        let decrypted1 = assert_ok!(encryption.decrypt(&encrypted1));
        let decrypted2 = assert_ok!(encryption.decrypt(&encrypted2));

        assert_eq!(decrypted1.expose_secret(), decrypted2.expose_secret());
        assert_eq!(decrypted1.expose_secret(), token);
    }

    #[test]
    fn test_invalid_encrypted_data() {
        let encryption = create_test_encryption();

        // Too short
        let err = assert_err!(encryption.decrypt(&[1, 2, 3]));
        assert_snapshot!(err, @"Invalid encrypted token: too short");

        // Invalid data
        let invalid_data = vec![0u8; 50];
        let err = assert_err!(encryption.decrypt(&invalid_data));
        assert_snapshot!(err, @"Failed to decrypt token");
    }

    #[test]
    fn test_different_keys() {
        let encryption1 = create_test_encryption();

        // Create a different encryption with a different key
        let cipher2 = Aes256Gcm::new_from_slice(b"different_key_32_bytes_long!!!!!").unwrap();
        let encryption2 = TokenEncryption { cipher: cipher2 };

        let token = "ghs_test_token_123456789";

        // Encrypt with encryption1
        let encrypted = assert_ok!(encryption1.encrypt(token));

        // Try to decrypt with encryption2 (should fail)
        let err = assert_err!(encryption2.decrypt(&encrypted));
        assert_snapshot!(err, @"Failed to decrypt token");

        // But encryption1 should still work
        let decrypted = assert_ok!(encryption1.decrypt(&encrypted));
        assert_eq!(decrypted.expose_secret(), token);
    }
}
