//! Audit Log Encryption Module
//!
//! Provides encryption for sensitive audit log fields

use aes_gcm::{
    aead::{rand_core::RngCore, Aead, KeyInit, OsRng},
    AeadCore, Aes256Gcm, Nonce,
};
use anyhow::{Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use serde::{Deserialize, Serialize};

/// Encrypted data wrapper
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedData {
    pub nonce: String,
    pub ciphertext: String,
}

/// Audit encryption manager
pub struct AuditEncryption {
    cipher: Aes256Gcm,
}

impl AuditEncryption {
    /// Create a new encryption manager with a key
    pub fn new(key: &[u8; 32]) -> Self {
        let cipher = Aes256Gcm::new_from_slice(key).expect("Failed to create cipher");
        Self { cipher }
    }

    /// Create from environment variable
    pub fn from_env() -> Result<Self> {
        let key_hex =
            std::env::var("AUDIT_ENCRYPTION_KEY").context("AUDIT_ENCRYPTION_KEY not set")?;

        let key_bytes = hex::decode(&key_hex).context("Invalid hex in AUDIT_ENCRYPTION_KEY")?;

        if key_bytes.len() != 32 {
            anyhow::bail!("AUDIT_ENCRYPTION_KEY must be 32 bytes (64 hex chars)");
        }

        let mut key = [0u8; 32];
        key.copy_from_slice(&key_bytes);
        Ok(Self::new(&key))
    }

    /// Generate a random key
    pub fn generate_key() -> [u8; 32] {
        let mut key = [0u8; 32];
        OsRng.fill_bytes(&mut key);
        key
    }

    /// Encrypt a string
    pub fn encrypt(&self, plaintext: &str) -> Result<EncryptedData> {
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let ciphertext = self
            .cipher
            .encrypt(&nonce, plaintext.as_bytes())
            .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;

        Ok(EncryptedData {
            nonce: BASE64.encode(nonce),
            ciphertext: BASE64.encode(ciphertext),
        })
    }

    /// Decrypt a string
    pub fn decrypt(&self, data: &EncryptedData) -> Result<String> {
        let nonce = BASE64.decode(&data.nonce).context("Invalid nonce base64")?;
        let ciphertext = BASE64
            .decode(&data.ciphertext)
            .context("Invalid ciphertext base64")?;

        let plaintext = self
            .cipher
            .decrypt(Nonce::from_slice(&nonce), ciphertext.as_ref())
            .map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))?;

        String::from_utf8(plaintext).context("Decrypted data is not valid UTF-8")
    }

    /// Encrypt sensitive audit fields
    pub fn encrypt_audit_entry(&self, entry: &AuditLogEntry) -> Result<EncryptedAuditLogEntry> {
        Ok(EncryptedAuditLogEntry {
            id: entry.id,
            key_id: entry.key_id.clone(),
            action: entry.action.clone(),
            actor: entry.actor.clone(),
            ip_address: entry
                .ip_address
                .as_ref()
                .map(|ip| self.encrypt(ip))
                .transpose()?,
            user_agent: entry
                .user_agent
                .as_ref()
                .map(|ua| self.encrypt(ua))
                .transpose()?,
            request_path: entry.request_path.clone(),
            response_code: entry.response_code,
            anomaly_type: entry.anomaly_type.clone(),
            details: entry
                .details
                .as_ref()
                .map(|d| self.encrypt(&d.to_string()))
                .transpose()?,
            created_at: entry.created_at,
        })
    }

    /// Decrypt sensitive audit fields
    pub fn decrypt_audit_entry(&self, entry: &EncryptedAuditLogEntry) -> Result<AuditLogEntry> {
        Ok(AuditLogEntry {
            id: entry.id,
            key_id: entry.key_id.clone(),
            action: entry.action.clone(),
            actor: entry.actor.clone(),
            ip_address: entry
                .ip_address
                .as_ref()
                .map(|enc| self.decrypt(enc))
                .transpose()?,
            user_agent: entry
                .user_agent
                .as_ref()
                .map(|enc| self.decrypt(enc))
                .transpose()?,
            request_path: entry.request_path.clone(),
            response_code: entry.response_code,
            anomaly_type: entry.anomaly_type.clone(),
            details: entry
                .details
                .as_ref()
                .map(|enc| self.decrypt(enc))
                .transpose()?
                .map(|s| serde_json::from_str(&s))
                .transpose()?,
            created_at: entry.created_at,
        })
    }
}

/// Audit log entry (plaintext)
#[derive(Debug, Clone)]
pub struct AuditLogEntry {
    pub id: i64,
    pub key_id: String,
    pub action: String,
    pub actor: Option<String>,
    pub ip_address: Option<String>,
    pub user_agent: Option<String>,
    pub request_path: Option<String>,
    pub response_code: Option<i32>,
    pub anomaly_type: Option<String>,
    pub details: Option<serde_json::Value>,
    pub created_at: chrono::DateTime<chrono::Utc>,
}

/// Audit log entry (with encrypted fields)
#[derive(Debug, Clone)]
pub struct EncryptedAuditLogEntry {
    pub id: i64,
    pub key_id: String,
    pub action: String,
    pub actor: Option<String>,
    pub ip_address: Option<EncryptedData>,
    pub user_agent: Option<EncryptedData>,
    pub request_path: Option<String>,
    pub response_code: Option<i32>,
    pub anomaly_type: Option<String>,
    pub details: Option<EncryptedData>,
    pub created_at: chrono::DateTime<chrono::Utc>,
}

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

    #[test]
    fn test_encrypt_decrypt() {
        let key = AuditEncryption::generate_key();
        let enc = AuditEncryption::new(&key);

        let plaintext = "sensitive data 12345";
        let encrypted = enc.encrypt(plaintext).unwrap();
        let decrypted = enc.decrypt(&encrypted).unwrap();

        assert_eq!(plaintext, decrypted);
    }

    #[test]
    fn test_different_nonces() {
        let key = AuditEncryption::generate_key();
        let enc = AuditEncryption::new(&key);

        let encrypted1 = enc.encrypt("same data").unwrap();
        let encrypted2 = enc.encrypt("same data").unwrap();

        // Different nonces should produce different ciphertexts
        assert_ne!(encrypted1.nonce, encrypted2.nonce);
        assert_ne!(encrypted1.ciphertext, encrypted2.ciphertext);

        // But both should decrypt to the same plaintext
        assert_eq!(
            enc.decrypt(&encrypted1).unwrap(),
            enc.decrypt(&encrypted2).unwrap()
        );
    }
}
