import os

from cryptography.fernet import Fernet, InvalidToken

ENCRYPTION_PREFIX = "enc:"
KEY_ENV_VAR = "BROKER_TOKEN_KEY"


def _get_fernet() -> Fernet:
    key = (os.getenv(KEY_ENV_VAR) or "").strip()
    if not key:
        raise RuntimeError(f"{KEY_ENV_VAR} is not set")
    try:
        return Fernet(key.encode("utf-8"))
    except Exception as exc:
        raise RuntimeError(
            f"{KEY_ENV_VAR} must be a urlsafe base64-encoded 32-byte key"
        ) from exc


def encrypt_value(value: str | None) -> str | None:
    if not value:
        return value
    if value.startswith(ENCRYPTION_PREFIX):
        return value
    token = _get_fernet().encrypt(value.encode("utf-8")).decode("utf-8")
    return f"{ENCRYPTION_PREFIX}{token}"


def decrypt_value(value: str | None) -> str | None:
    if not value:
        return value
    if not value.startswith(ENCRYPTION_PREFIX):
        return value
    token = value[len(ENCRYPTION_PREFIX) :]
    try:
        return _get_fernet().decrypt(token.encode("utf-8")).decode("utf-8")
    except InvalidToken as exc:
        raise RuntimeError("Unable to decrypt token; invalid BROKER_TOKEN_KEY") from exc