from cryptography.fernet import Fernet
from flask import Flask


class SimpleEncryption:
    def __init__(self, app: Flask = None):
        self.app = app
        self.fernet = None
        if app is not None:
            self.init_app(app)

    def init_app(self, app: Flask):
        self.app = app
        encryption_key = app.config.get('API_ENCRYPTION_KEY')
        if not encryption_key:
            raise ValueError("ENCRYPTION_KEY is not set in the app configuration")
        self.fernet = Fernet(encryption_key.encode())

        # Optionally log the initialization (similar to your current setup)
        app.logger.info('SimpleEncryption initialized')

    def encrypt_api_key(self, api_key: str) -> str:
        if not self.fernet:
            raise RuntimeError("SimpleEncryption is not initialized. Call init_app first.")
        return self.fernet.encrypt(api_key.encode()).decode()

    def decrypt_api_key(self, encrypted_api_key: str) -> str:
        if not self.fernet:
            raise RuntimeError("SimpleEncryption is not initialized. Call init_app first.")
        return self.fernet.decrypt(encrypted_api_key.encode()).decode()

    @staticmethod
    def generate_key() -> str:
        """Generate a new Fernet key."""
        return Fernet.generate_key().decode()

# Usage:
# from common.utils.simple_encryption import simple_encryption
# simple_encryption.init_app(app)