from cryptography.fernet import Fernet
import os

DEFAULT_KEY = "fdV07LMuahTZu_o4T311JTJkR9UaagbxWVcE-vSScwA="

def get_key() -> bytes:
    key = os.getenv("MATH_CALC_KEY", DEFAULT_KEY)
    return key.encode() if isinstance(key, str) else key

def encrypt_message(message: str) -> str:
    key = get_key()
    f = Fernet(key)
    encrypted = f.encrypt(message.encode('utf-8'))
    return encrypted.decode('utf-8')

def decrypt_message(token: str) -> str:
    key = get_key()
    f = Fernet(key)
    decrypted = f.decrypt(token.encode('utf-8'))
    return decrypted.decode('utf-8')

def generate_key() -> str:
    return Fernet.generate_key().decode('utf-8')
