"""In-memory token cache for aablocks.

Shared between auth and logging modules. Tokens are keyed by auth_domain
so users can be logged in to multiple environments simultaneously.
"""

from __future__ import annotations

import threading
from typing import Any

_cached_tokens: dict[str, dict[str, Any]] = {}
_token_lock = threading.Lock()


def get_cached_token(auth_domain: str) -> dict[str, Any] | None:
    """Get the token for an auth domain from the in-memory cache."""
    with _token_lock:
        return _cached_tokens.get(auth_domain)


def set_cached_token(auth_domain: str, token: dict[str, Any]) -> None:
    """Store the token for an auth domain in the in-memory cache."""
    with _token_lock:
        _cached_tokens[auth_domain] = token


def clear_cached_token(auth_domain: str) -> None:
    """Clear the token for an auth domain from the in-memory cache."""
    with _token_lock:
        _cached_tokens.pop(auth_domain, None)
