foxhunt/services/ml_training_service/src/encryption.rs

//! Model Encryption Key Management with Vault Integration
//!
//! This module provides secure encryption key management for ML model storage,
//! supporting key rotation, multiple algorithms, and secure key retrieval from Vault.

use aes_gcm::aead::{Aead, Payload};
use aes_gcm::{Aes256Gcm, KeyInit, Nonce as AesNonce};
use chacha20poly1305::{ChaCha20Poly1305, Nonce as ChaChaNonce};
use pbkdf2::pbkdf2_hmac;
use sha2::Sha256;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use zeroize::Zeroize;

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tokio::fs;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

use config::structures::EncryptionConfig;

/// Encryption keys structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionKeys {
    pub primary_key: String,
    pub key_id: String,
    pub algorithm: String,
    pub created_at: SystemTime,
}

impl EncryptionKeys {
    /// Check if key should be rotated based on age
    pub fn should_rotate(&self, rotation_days: u64) -> bool {
        match self.created_at.elapsed() {
            Ok(elapsed) => elapsed.as_secs() > rotation_days * 86400,
            Err(_) => true, // If we can't determine age, assume rotation needed
        }
    }
}

/// Encryption key manager with secure configuration
pub struct EncryptionKeyManager {
    config: EncryptionConfig,
    config_loader: Option<Arc<config::database::PostgresConfigLoader>>,
    cached_keys: Arc<RwLock<Option<CachedEncryptionKeys>>>,
}

/// Cached encryption keys with metadata
#[derive(Debug, Clone)]
struct CachedEncryptionKeys {
    keys: EncryptionKeys,
    cached_at: SystemTime,
    cache_ttl_secs: u64,
}

impl CachedEncryptionKeys {
    fn new(keys: EncryptionKeys, cache_ttl_secs: u64) -> Self {
        Self {
            keys,
            cached_at: SystemTime::now(),
            cache_ttl_secs,
        }
    }

    fn is_expired(&self) -> bool {
        match self.cached_at.elapsed() {
            Ok(elapsed) => elapsed.as_secs() > self.cache_ttl_secs,
            Err(_) => true, // If we can't determine time, assume expired
        }
    }

    #[allow(dead_code)]
    fn needs_rotation(&self, rotation_days: u64) -> bool {
        self.keys.should_rotate(rotation_days)
    }
}

/// Encryption algorithm configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionAlgorithmConfig {
    pub algorithm: EncryptionAlgorithm,
    pub key_size_bits: usize,
    pub iv_size_bytes: usize,
    pub tag_size_bytes: usize,
}

/// Supported encryption algorithms
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum EncryptionAlgorithm {
    #[serde(rename = "AES-256-GCM")]
    Aes256Gcm,
    #[serde(rename = "ChaCha20Poly1305")]
    ChaCha20Poly1305,
    #[serde(rename = "AES-256-CTR")]
    Aes256Ctr,
}

impl EncryptionAlgorithm {
    /// Get algorithm configuration
    #[allow(dead_code)]
    pub fn get_config(&self) -> EncryptionAlgorithmConfig {
        match self {
            Self::Aes256Gcm => EncryptionAlgorithmConfig {
                algorithm: self.clone(),
                key_size_bits: 256,
                iv_size_bytes: 12,  // 96-bit IV for GCM
                tag_size_bytes: 16, // 128-bit authentication tag
            },
            Self::ChaCha20Poly1305 => EncryptionAlgorithmConfig {
                algorithm: self.clone(),
                key_size_bits: 256,
                iv_size_bytes: 12,  // 96-bit nonce
                tag_size_bytes: 16, // 128-bit authentication tag
            },
            Self::Aes256Ctr => EncryptionAlgorithmConfig {
                algorithm: self.clone(),
                key_size_bits: 256,
                iv_size_bytes: 16, // 128-bit IV for CTR
                tag_size_bytes: 0, // No authentication tag (not AEAD)
            },
        }
    }

    /// Check if algorithm provides authenticated encryption
    #[allow(dead_code)]
    pub fn is_authenticated(&self) -> bool {
        matches!(self, Self::Aes256Gcm | Self::ChaCha20Poly1305)
    }
}

impl std::str::FromStr for EncryptionAlgorithm {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self> {
        match s {
            "AES-256-GCM" => Ok(Self::Aes256Gcm),
            "ChaCha20Poly1305" => Ok(Self::ChaCha20Poly1305),
            "AES-256-CTR" => Ok(Self::Aes256Ctr),
            _ => Err(anyhow::anyhow!("Unsupported encryption algorithm: {}", s)),
        }
    }
}

impl std::fmt::Display for EncryptionAlgorithm {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let name = match self {
            Self::Aes256Gcm => "AES-256-GCM",
            Self::ChaCha20Poly1305 => "ChaCha20Poly1305",
            Self::Aes256Ctr => "AES-256-CTR",
        };
        write!(f, "{}", name)
    }
}

/// Encryption metadata for stored models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptionMetadata {
    pub algorithm: EncryptionAlgorithm,
    pub key_id: String,
    pub nonce: Vec<u8>,       // 96-bit nonce (12 bytes)
    pub salt: Vec<u8>,        // Salt for key derivation (16 bytes)
    pub tag: Option<Vec<u8>>, // For AEAD algorithms
    pub encrypted_at: SystemTime,
    pub key_version: u32,
}

#[allow(dead_code)]
impl EncryptionKeyManager {
    /// Create a new encryption key manager
    pub fn new(
        config: EncryptionConfig,
        config_loader: Option<Arc<config::database::PostgresConfigLoader>>,
    ) -> Self {
        Self {
            config,
            config_loader,
            cached_keys: Arc::new(RwLock::new(None)),
        }
    }

    /// Check if encryption is enabled
    pub fn is_encryption_enabled(&self) -> bool {
        self.config.enable_encryption
    }

    /// Get current encryption algorithm
    pub fn get_algorithm(&self) -> Result<EncryptionAlgorithm> {
        self.config.algorithm.parse()
    }

    /// Load encryption keys from secure configuration or fallback source
    pub async fn load_encryption_keys(&self) -> Result<EncryptionKeys> {
        // Check cache first
        {
            let cached_guard = self.cached_keys.read().await;
            if let Some(cached) = cached_guard.as_ref() {
                if !cached.is_expired() {
                    debug!("Using cached encryption keys");
                    return Ok(cached.keys.clone());
                }
            }
        }

        // Load from fallback source (local file or generated)
        // Note: PostgresConfigLoader doesn't provide encryption key storage
        // Keys should be managed through local files or generated securely
        let keys = if self.config_loader.is_some() {
            debug!("Config loader available but encryption keys managed locally");
            self.load_fallback_keys().await?
        } else {
            info!("Loading encryption keys from fallback source (no secure configuration)");
            self.load_fallback_keys().await?
        };

        // Cache the loaded keys
        {
            let mut cached_guard = self.cached_keys.write().await;
            *cached_guard = Some(CachedEncryptionKeys::new(keys.clone(), 300)); // 5-minute cache
        }

        debug!("Encryption keys loaded and cached");
        Ok(keys)
    }

    /// Load encryption keys from fallback source (local file or generated)
    async fn load_fallback_keys(&self) -> Result<EncryptionKeys> {
        if let Some(key_file) = &self.config.local_key_file {
            let path = PathBuf::from(key_file);
            self.load_keys_from_file(&path).await
        } else {
            warn!("No fallback key source configured, generating temporary keys");
            self.generate_temporary_keys().await
        }
    }

    /// Generate temporary encryption keys (wrapper for generate_secure_keys)
    async fn generate_temporary_keys(&self) -> Result<EncryptionKeys> {
        self.generate_secure_keys().await
    }

    /// Load encryption keys from local file
    async fn load_keys_from_file(&self, key_file: &PathBuf) -> Result<EncryptionKeys> {
        let key_data = fs::read_to_string(key_file)
            .await
            .context("Failed to read encryption key file")?;

        let keys: EncryptionKeys =
            serde_json::from_str(&key_data).context("Failed to parse encryption key file")?;

        info!("Loaded encryption keys from file: {}", key_file.display());
        Ok(keys)
    }

    /// Derive encryption key from base key using PBKDF2
    fn derive_key(&self, base_key: &str, salt: &[u8]) -> Result<[u8; 32]> {
        let mut derived_key = [0u8; 32];

        // Use PBKDF2 with 100,000 iterations (NIST recommendation)
        pbkdf2_hmac::<Sha256>(base_key.as_bytes(), salt, 100_000, &mut derived_key);

        Ok(derived_key)
    }

    /// Generate cryptographically secure random nonce
    fn generate_nonce(&self, size: usize) -> Result<Vec<u8>> {
        use rand::{rngs::OsRng, RngCore};

        let mut nonce = vec![0u8; size];
        OsRng.fill_bytes(&mut nonce);

        Ok(nonce)
    }

    /// Generate salt for key derivation
    fn generate_salt(&self) -> Result<[u8; 16]> {
        use rand::{rngs::OsRng, RngCore};
        let mut salt = [0u8; 16];
        OsRng.fill_bytes(&mut salt);
        Ok(salt)
    }

    /// Generate cryptographically secure encryption keys
    async fn generate_secure_keys(&self) -> Result<EncryptionKeys> {
        info!("Generating cryptographically secure encryption keys using OsRng");

        // Use cryptographically secure random number generator
        use base64::Engine;
        use rand::{rngs::OsRng, RngCore};
        let mut key_bytes = vec![0u8; 32];
        OsRng.fill_bytes(&mut key_bytes);
        let primary_key = base64::prelude::BASE64_STANDARD.encode(&key_bytes);

        let keys = EncryptionKeys {
            primary_key,
            key_id: format!(
                "temp-key-{}",
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs()
            ),
            algorithm: self.config.algorithm.clone(),
            created_at: SystemTime::now(),
        };

        Ok(keys)
    }

    /// Check if current keys need rotation
    pub async fn keys_need_rotation(&self) -> Result<bool> {
        let keys = self.load_encryption_keys().await?;
        Ok(keys.should_rotate(self.config.key_rotation_days))
    }

    /// Request key rotation (would trigger Vault key rotation in production)
    pub async fn request_key_rotation(&self) -> Result<()> {
        info!("Requesting encryption key rotation");

        // Clear cache to force reload of new keys
        {
            let mut cached_guard = self.cached_keys.write().await;
            *cached_guard = None;
        }

        // In a production environment, this would:
        // 1. Request new keys from Vault
        // 2. Update Vault secrets with new keys
        // 3. Notify other services of key rotation
        // 4. Schedule old key deprecation

        warn!("Key rotation requested - implement Vault key rotation logic for production");
        Ok(())
    }

    /// Encrypt model data
    pub async fn encrypt_model_data(&self, data: &[u8]) -> Result<(Vec<u8>, EncryptionMetadata)> {
        if !self.config.enable_encryption {
            return Err(anyhow::anyhow!("Encryption is disabled"));
        }

        let keys = self.load_encryption_keys().await?;
        let algorithm = self.get_algorithm()?;

        // Generate salt for key derivation
        let salt = self.generate_salt()?;

        // Generate random IV/nonce
        let nonce = self.generate_nonce(12)?; // 96-bit nonce for GCM/ChaCha20-Poly1305

        // Perform authenticated encryption
        let (encrypted_data, tag) =
            self.perform_encryption(data, &keys.primary_key, &nonce, &salt, &algorithm)?;

        let metadata = EncryptionMetadata {
            algorithm,
            key_id: keys.key_id,
            nonce,
            salt: salt.to_vec(),
            tag: Some(tag), // Authentication tag from AEAD
            encrypted_at: SystemTime::now(),
            key_version: 1, // Would track actual key versions
        };

        debug!(
            "Encrypted {} bytes of model data with algorithm {}",
            data.len(),
            metadata.algorithm
        );

        Ok((encrypted_data, metadata))
    }

    /// Decrypt model data
    pub async fn decrypt_model_data(
        &self,
        encrypted_data: &[u8],
        metadata: &EncryptionMetadata,
    ) -> Result<Vec<u8>> {
        if !self.config.enable_encryption {
            return Err(anyhow::anyhow!("Encryption is disabled"));
        }

        let keys = self.load_encryption_keys().await?;

        // Verify key ID matches (in production, support multiple key versions)
        if keys.key_id != metadata.key_id {
            return Err(anyhow::anyhow!(
                "Key ID mismatch: expected {}, got {}",
                keys.key_id,
                metadata.key_id
            ));
        }

        let decrypted_data = self.perform_decryption(
            encrypted_data,
            &keys.primary_key,
            &metadata.nonce,
            &metadata.salt,
            metadata.tag.as_deref(),
            &metadata.algorithm,
        )?;

        debug!(
            "Decrypted {} bytes of model data with algorithm {}",
            decrypted_data.len(),
            metadata.algorithm
        );

        Ok(decrypted_data)
    }

    /// Perform actual encryption (simplified implementation)
    fn perform_encryption(
        &self,
        data: &[u8],
        key: &str,
        iv: &[u8],
        salt: &[u8],
        algorithm: &EncryptionAlgorithm,
    ) -> Result<(Vec<u8>, Vec<u8>)> {
        match algorithm {
            EncryptionAlgorithm::Aes256Gcm => {
                // Use AES-256-GCM encryption
                self.aes_gcm_encrypt(data, key, iv, salt)
            },
            EncryptionAlgorithm::ChaCha20Poly1305 => {
                // Use ChaCha20-Poly1305 encryption
                self.chacha20_encrypt(data, key, iv, salt)
            },
            EncryptionAlgorithm::Aes256Ctr => Err(anyhow::anyhow!(
                "AES-256-CTR is not authenticated, use AES-256-GCM instead"
            )),
        }
    }

    /// Perform actual decryption (simplified implementation)
    fn perform_decryption(
        &self,
        encrypted_data: &[u8],
        key: &str,
        iv: &[u8],
        salt: &[u8],
        tag: Option<&[u8]>,
        algorithm: &EncryptionAlgorithm,
    ) -> Result<Vec<u8>> {
        let tag = tag.ok_or_else(|| anyhow::anyhow!("Missing authentication tag for AEAD"))?;

        match algorithm {
            EncryptionAlgorithm::Aes256Gcm => {
                // Use AES-256-GCM decryption
                self.aes_gcm_decrypt(encrypted_data, key, iv, salt, tag)
            },
            EncryptionAlgorithm::ChaCha20Poly1305 => {
                // Use ChaCha20-Poly1305 decryption
                self.chacha20_decrypt(encrypted_data, key, iv, salt, tag)
            },
            EncryptionAlgorithm::Aes256Ctr => Err(anyhow::anyhow!(
                "AES-256-CTR is not authenticated, use AES-256-GCM instead"
            )),
        }
    }

    // Production AES-256-GCM encryption
    fn aes_gcm_encrypt(
        &self,
        data: &[u8],
        base_key: &str,
        nonce: &[u8],
        salt: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>)> {
        // Derive key using PBKDF2
        let mut derived_key = self.derive_key(base_key, salt)?;

        // Create cipher
        let cipher = Aes256Gcm::new_from_slice(&derived_key)
            .map_err(|e| anyhow::anyhow!("Failed to create AES-256-GCM cipher: {}", e))?;

        // Create nonce (96 bits = 12 bytes)
        let nonce_array = AesNonce::from_slice(nonce);

        // Encrypt with additional authenticated data
        let ciphertext = cipher
            .encrypt(
                nonce_array,
                Payload {
                    msg: data,
                    aad: b"foxhunt-ml-model-v1", // Additional authenticated data
                },
            )
            .map_err(|e| anyhow::anyhow!("AES-256-GCM encryption failed: {}", e))?;

        // Zero out derived key
        derived_key.zeroize();

        // Split ciphertext and tag (tag is last 16 bytes)
        let tag_offset = ciphertext.len() - 16;
        let encrypted_data = ciphertext[..tag_offset].to_vec();
        let tag = ciphertext[tag_offset..].to_vec();

        info!(
            "Successfully encrypted {} bytes with AES-256-GCM",
            data.len()
        );
        Ok((encrypted_data, tag))
    }

    fn aes_gcm_decrypt(
        &self,
        encrypted_data: &[u8],
        base_key: &str,
        nonce: &[u8],
        salt: &[u8],
        tag: &[u8],
    ) -> Result<Vec<u8>> {
        // Derive key using PBKDF2
        let mut derived_key = self.derive_key(base_key, salt)?;

        // Create cipher
        let cipher = Aes256Gcm::new_from_slice(&derived_key)
            .map_err(|e| anyhow::anyhow!("Failed to create AES-256-GCM cipher: {}", e))?;

        // Create nonce
        let nonce_array = AesNonce::from_slice(nonce);

        // Combine ciphertext and tag
        let mut ciphertext_with_tag = encrypted_data.to_vec();
        ciphertext_with_tag.extend_from_slice(tag);

        // Decrypt with AAD verification
        let plaintext = cipher
            .decrypt(
                nonce_array,
                Payload {
                    msg: &ciphertext_with_tag,
                    aad: b"foxhunt-ml-model-v1",
                },
            )
            .map_err(|e| {
                anyhow::anyhow!(
                    "AES-256-GCM decryption failed (authentication error): {}",
                    e
                )
            })?;

        // Zero out derived key
        derived_key.zeroize();

        info!(
            "Successfully decrypted {} bytes with AES-256-GCM",
            plaintext.len()
        );
        Ok(plaintext)
    }

    fn chacha20_encrypt(
        &self,
        data: &[u8],
        base_key: &str,
        nonce: &[u8],
        salt: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>)> {
        // Derive key using PBKDF2
        let mut derived_key = self.derive_key(base_key, salt)?;

        // Create cipher
        let cipher = ChaCha20Poly1305::new_from_slice(&derived_key)
            .map_err(|e| anyhow::anyhow!("Failed to create ChaCha20-Poly1305 cipher: {}", e))?;

        // Create nonce (96 bits = 12 bytes)
        let nonce_array = ChaChaNonce::from_slice(nonce);

        // Encrypt with AAD
        let ciphertext = cipher
            .encrypt(
                nonce_array,
                Payload {
                    msg: data,
                    aad: b"foxhunt-ml-model-v1",
                },
            )
            .map_err(|e| anyhow::anyhow!("ChaCha20-Poly1305 encryption failed: {}", e))?;

        // Zero out derived key
        derived_key.zeroize();

        // Split ciphertext and tag
        let tag_offset = ciphertext.len() - 16;
        let encrypted_data = ciphertext[..tag_offset].to_vec();
        let tag = ciphertext[tag_offset..].to_vec();

        info!(
            "Successfully encrypted {} bytes with ChaCha20-Poly1305",
            data.len()
        );
        Ok((encrypted_data, tag))
    }

    fn chacha20_decrypt(
        &self,
        encrypted_data: &[u8],
        base_key: &str,
        nonce: &[u8],
        salt: &[u8],
        tag: &[u8],
    ) -> Result<Vec<u8>> {
        // Derive key using PBKDF2
        let mut derived_key = self.derive_key(base_key, salt)?;

        // Create cipher
        let cipher = ChaCha20Poly1305::new_from_slice(&derived_key)
            .map_err(|e| anyhow::anyhow!("Failed to create ChaCha20-Poly1305 cipher: {}", e))?;

        // Create nonce
        let nonce_array = ChaChaNonce::from_slice(nonce);

        // Combine ciphertext and tag
        let mut ciphertext_with_tag = encrypted_data.to_vec();
        ciphertext_with_tag.extend_from_slice(tag);

        // Decrypt with AAD verification
        let plaintext = cipher
            .decrypt(
                nonce_array,
                Payload {
                    msg: &ciphertext_with_tag,
                    aad: b"foxhunt-ml-model-v1",
                },
            )
            .map_err(|e| {
                anyhow::anyhow!(
                    "ChaCha20-Poly1305 decryption failed (authentication error): {}",
                    e
                )
            })?;

        // Zero out derived key
        derived_key.zeroize();

        info!(
            "Successfully decrypted {} bytes with ChaCha20-Poly1305",
            plaintext.len()
        );
        Ok(plaintext)
    }

    /// Get encryption statistics
    pub async fn get_encryption_stats(&self) -> Result<EncryptionStats> {
        let keys = self.load_encryption_keys().await?;
        let needs_rotation = keys.should_rotate(self.config.key_rotation_days);

        Ok(EncryptionStats {
            encryption_enabled: self.config.enable_encryption,
            algorithm: self.get_algorithm()?,
            key_id: keys.key_id,
            key_age_days: keys
                .created_at
                .elapsed()
                .map(|d| d.as_secs() / 86400)
                .unwrap_or(0),
            needs_rotation,
            rotation_interval_days: self.config.key_rotation_days,
            cache_hit_rate: 0.0, // Would track actual cache statistics
        })
    }
}

/// Encryption statistics
#[allow(dead_code)]
#[derive(Debug, Clone, Serialize)]
pub struct EncryptionStats {
    pub encryption_enabled: bool,
    pub algorithm: EncryptionAlgorithm,
    pub key_id: String,
    pub key_age_days: u64,
    pub needs_rotation: bool,
    pub rotation_interval_days: u64,
    pub cache_hit_rate: f64,
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_encryption_algorithm_parsing() {
        assert_eq!(
            "AES-256-GCM".parse::<EncryptionAlgorithm>().unwrap(),
            EncryptionAlgorithm::Aes256Gcm
        );
        assert_eq!(
            "ChaCha20Poly1305".parse::<EncryptionAlgorithm>().unwrap(),
            EncryptionAlgorithm::ChaCha20Poly1305
        );
        assert!("Invalid-Algorithm".parse::<EncryptionAlgorithm>().is_err());
    }

    #[test]
    fn test_algorithm_config() {
        let aes_config = EncryptionAlgorithm::Aes256Gcm.get_config();
        assert_eq!(aes_config.key_size_bits, 256);
        assert_eq!(aes_config.iv_size_bytes, 12);
        assert!(aes_config.algorithm.is_authenticated());

        let ctr_config = EncryptionAlgorithm::Aes256Ctr.get_config();
        assert!(!ctr_config.algorithm.is_authenticated());
    }

    #[tokio::test]
    async fn test_encryption_key_manager_creation() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        assert!(manager.is_encryption_enabled());
        assert_eq!(
            manager.get_algorithm().unwrap(),
            EncryptionAlgorithm::Aes256Gcm
        );
    }

    #[tokio::test]
    async fn test_temporary_key_generation() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let keys = manager.generate_temporary_keys().await.unwrap();

        assert!(!keys.primary_key.is_empty());
        assert!(!keys.key_id.is_empty());
        assert!(keys.key_id.starts_with("temp-key-"));
    }

    #[tokio::test]
    async fn test_aes_gcm_encryption_decryption() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let test_data = b"Hello, encrypted world!";

        let (encrypted, metadata) = manager.encrypt_model_data(test_data).await.unwrap();

        // Verify encrypted data is different from plaintext
        assert_ne!(encrypted.as_slice(), test_data);

        // Verify metadata
        assert_eq!(metadata.algorithm, EncryptionAlgorithm::Aes256Gcm);
        assert_eq!(metadata.nonce.len(), 12); // 96-bit nonce
        assert_eq!(metadata.salt.len(), 16); // 128-bit salt
        assert!(metadata.tag.is_some());
        assert_eq!(metadata.tag.as_ref().expect("INVARIANT: Option should be Some").len(), 16); // 128-bit tag

        // Verify decryption works
        let decrypted = manager
            .decrypt_model_data(&encrypted, &metadata)
            .await
            .unwrap();
        assert_eq!(decrypted, test_data);
    }

    #[tokio::test]
    async fn test_chacha20_encryption_decryption() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "ChaCha20Poly1305".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let test_data = b"Testing ChaCha20-Poly1305 encryption!";

        let (encrypted, metadata) = manager.encrypt_model_data(test_data).await.unwrap();

        // Verify encrypted data is different
        assert_ne!(encrypted.as_slice(), test_data);

        // Verify metadata
        assert_eq!(metadata.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
        assert!(metadata.tag.is_some());

        // Verify decryption
        let decrypted = manager
            .decrypt_model_data(&encrypted, &metadata)
            .await
            .unwrap();
        assert_eq!(decrypted, test_data);
    }

    #[tokio::test]
    async fn test_encryption_authentication_tag_validation() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let test_data = b"Testing authentication tag validation";

        let (encrypted, mut metadata) = manager.encrypt_model_data(test_data).await.unwrap();

        // Tamper with the tag
        if let Some(ref mut tag) = metadata.tag {
            if let Some(first_byte) = tag.first_mut() {
                *first_byte ^= 0xFF; // Flip bits in first byte
            }
        }

        // Decryption should fail with tampered tag
        let result = manager.decrypt_model_data(&encrypted, &metadata).await;
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("authentication error"));
    }

    #[tokio::test]
    async fn test_nonce_uniqueness() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);
        let test_data = b"Testing nonce uniqueness";

        // Encrypt same data multiple times
        let (_, metadata1) = manager.encrypt_model_data(test_data).await.unwrap();
        let (_, metadata2) = manager.encrypt_model_data(test_data).await.unwrap();
        let (_, metadata3) = manager.encrypt_model_data(test_data).await.unwrap();

        // Nonces should be different (cryptographically random)
        assert_ne!(metadata1.nonce, metadata2.nonce);
        assert_ne!(metadata1.nonce, metadata3.nonce);
        assert_ne!(metadata2.nonce, metadata3.nonce);
    }

    #[tokio::test]
    async fn test_large_data_encryption() {
        let config = EncryptionConfig {
            enable_encryption: true,
            algorithm: "AES-256-GCM".to_string(),
            key_rotation_days: 30,
            encryption_keys_vault_path: None,
            local_key_file: None,
        };

        let manager = EncryptionKeyManager::new(config, None);

        // Test with 1MB of data
        let large_data = vec![0xAB; 1024 * 1024];

        let (encrypted, metadata) = manager.encrypt_model_data(&large_data).await.unwrap();

        // Verify encryption
        assert_ne!(encrypted.as_slice(), large_data.as_slice());

        // Verify decryption
        let decrypted = manager
            .decrypt_model_data(&encrypted, &metadata)
            .await
            .unwrap();
        assert_eq!(decrypted, large_data);
    }
}