"""
AES-256 Encryption Module for API Keys
Phase 3: Security Enhancement
"""

import os
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from typing import Optional
import secrets

class APIKeyEncryption:
    def __init__(self, master_password: Optional[str] = None):
        """
        Initialize encryption with master password
        If no password provided, uses environment variable or generates one
        """
        self.master_password = master_password or os.getenv("ENCRYPTION_KEY") or self._generate_master_key()
        self.salt = b'ai_server_salt_2025'  # Fixed salt for consistency
        self._fernet = self._create_fernet()
    
    def _generate_master_key(self) -> str:
        """Generate a secure master key"""
        key = secrets.token_urlsafe(32)
        print(f"🔑 Generated new encryption key: {key}")
        print("⚠️  IMPORTANT: Save this key in your environment variables!")
        print(f"   export ENCRYPTION_KEY='{key}'")
        return key
    
    def _create_fernet(self) -> Fernet:
        """Create Fernet cipher from master password"""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=self.salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(self.master_password.encode()))
        return Fernet(key)
    
    def encrypt_api_key(self, api_key: str) -> str:
        """Encrypt API key using AES-256"""
        try:
            encrypted_bytes = self._fernet.encrypt(api_key.encode())
            return base64.urlsafe_b64encode(encrypted_bytes).decode()
        except Exception as e:
            raise Exception(f"Encryption failed: {str(e)}")
    
    def decrypt_api_key(self, encrypted_api_key: str) -> str:
        """Decrypt API key"""
        try:
            encrypted_bytes = base64.urlsafe_b64decode(encrypted_api_key.encode())
            decrypted_bytes = self._fernet.decrypt(encrypted_bytes)
            return decrypted_bytes.decode()
        except Exception as e:
            raise Exception(f"Decryption failed: {str(e)}")
    
    def is_encrypted(self, api_key: str) -> bool:
        """Check if API key is already encrypted"""
        try:
            # Try to decode as base64 - encrypted keys are base64 encoded
            base64.urlsafe_b64decode(api_key.encode())
            # Try to decrypt - if successful, it's encrypted
            self.decrypt_api_key(api_key)
            return True
        except:
            return False
    
    def rotate_encryption_key(self, new_master_password: str, encrypted_keys: list) -> list:
        """Rotate encryption key - decrypt with old key, encrypt with new key"""
        old_fernet = self._fernet
        
        # Create new Fernet with new password
        self.master_password = new_master_password
        self._fernet = self._create_fernet()
        
        rotated_keys = []
        for encrypted_key in encrypted_keys:
            try:
                # Decrypt with old key
                decrypted = old_fernet.decrypt(base64.urlsafe_b64decode(encrypted_key.encode()))
                # Encrypt with new key
                new_encrypted = self.encrypt_api_key(decrypted.decode())
                rotated_keys.append(new_encrypted)
            except Exception as e:
                print(f"Failed to rotate key: {e}")
                rotated_keys.append(encrypted_key)  # Keep original if rotation fails
        
        return rotated_keys

# Global encryption instance
encryption = APIKeyEncryption()

def encrypt_api_key(api_key: str) -> str:
    """Convenience function to encrypt API key"""
    return encryption.encrypt_api_key(api_key)

def decrypt_api_key(encrypted_api_key: str) -> str:
    """Convenience function to decrypt API key"""
    return encryption.decrypt_api_key(encrypted_api_key)

def is_encrypted(api_key: str) -> bool:
    """Convenience function to check if API key is encrypted"""
    return encryption.is_encrypted(api_key)

# Test the encryption system
if __name__ == "__main__":
    # Test encryption/decryption
    test_key = "test-api-key-12345"
    
    print("🧪 Testing API Key Encryption...")
    print(f"Original: {test_key}")
    
    # Encrypt
    encrypted = encrypt_api_key(test_key)
    print(f"Encrypted: {encrypted}")
    
    # Decrypt
    decrypted = decrypt_api_key(encrypted)
    print(f"Decrypted: {decrypted}")
    
    # Verify
    print(f"✅ Match: {test_key == decrypted}")
    print(f"🔒 Is Encrypted: {is_encrypted(encrypted)}")
    print(f"🔓 Is Plain: {not is_encrypted(test_key)}")