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


def get_encryption_key() -> bytes:
    """Get or generate encryption key from environment variable"""
    key_string = os.environ.get('ENCRYPTION_KEY', 'default-encryption-key-change-in-production')
    
    # Derive a proper Fernet key from the string
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=b'aws-scanner-salt',  # In production, use a proper random salt
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(key_string.encode()))
    return key


def encrypt_value(value: str) -> str:
    """Encrypt a string value"""
    if not value:
        return None
    
    key = get_encryption_key()
    f = Fernet(key)
    encrypted = f.encrypt(value.encode('utf-8'))
    return base64.urlsafe_b64encode(encrypted).decode('utf-8')


def decrypt_value(encrypted_value: str) -> str:
    """Decrypt an encrypted string value"""
    if not encrypted_value:
        return None
    
    key = get_encryption_key()
    f = Fernet(key)
    encrypted_bytes = base64.urlsafe_b64decode(encrypted_value.encode('utf-8'))
    decrypted = f.decrypt(encrypted_bytes)
    return decrypted.decode('utf-8')