import os
import json
from cryptography.fernet import Fernet
from dotenv import load_dotenv

# Load environment variables from a .env file for local development
load_dotenv()

def _get_fernet_instance() -> Fernet:
    """
    Helper function to get the Fernet instance.
    This ensures the key is checked only when encryption/decryption is needed.
    """
    # It is CRITICAL that this key is set in your environment and kept secret.
    # It should be a 32-url-safe-base64-encoded key.
    encryption_key = os.getenv("ENCRYPTION_KEY")

    if not encryption_key:
        raise ValueError("ENCRYPTION_KEY is not set in the environment. Please generate a key and add it to your .env file.")

    # Ensure the key is in bytes for the Fernet instance
    return Fernet(encryption_key.encode())

def encrypt_tokens(access_token: str, refresh_token: str) -> str:
    """
    Combines access and refresh tokens into a JSON object, then encrypts it.
    """
    fernet = _get_fernet_instance()
    tokens = {"access_token": access_token, "refresh_token": refresh_token}
    tokens_json_string = json.dumps(tokens)
    encrypted_data = fernet.encrypt(tokens_json_string.encode())
    return encrypted_data.decode()

def decrypt_tokens(encrypted_data_str: str) -> dict:
    """
    Decrypts the token string back into a dictionary of tokens.
    """
    fernet = _get_fernet_instance()
    decrypted_data_bytes = fernet.decrypt(encrypted_data_str.encode())
    tokens_json_string = decrypted_data_bytes.decode()
    return json.loads(tokens_json_string)

def generate_key():
    """
    Utility function to generate a new encryption key. Run this once.
    """
    return Fernet.generate_key().decode()