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()

# 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
fernet = 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.
    """
    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.
    """
    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()