"""Keycloak authentication helpers for the Andlake platform.

Fetches access tokens via the OAuth2 client-credentials grant so that
SDK functions (MLflow, Nessie, etc.) can authenticate with platform
services without requiring a pre-injected token.
"""

from __future__ import annotations

import time
from urllib.request import Request, urlopen
from urllib.parse import urlencode
import json

from andlake.config import get_config

# Module-level token cache
_cached_token: str | None = None
_token_expires_at: float = 0.0


def get_access_token(
    *,
    client_id: str | None = None,
    client_secret: str | None = None,
    keycloak_url: str | None = None,
) -> str:
    """Return a valid Keycloak access token, fetching one if needed.

    Uses the OAuth2 **client credentials** grant.  Tokens are cached
    until 60 seconds before expiry.

    Parameters
    ----------
    client_id:
        Keycloak client ID.  Falls back to ``ANDLAKE_KEYCLOAK_CLIENT_ID``.
    client_secret:
        Keycloak client secret.  Falls back to ``ANDLAKE_KEYCLOAK_CLIENT_SECRET``.
    keycloak_url:
        Keycloak realm URL (e.g. ``https://keycloak.andlab.ai/realms/andlake``).
        Falls back to ``ANDLAKE_KEYCLOAK_URL``.

    Returns
    -------
    str
        A Bearer access token (JWT).

    Raises
    ------
    RuntimeError
        If client credentials are not configured or the token request fails.
    """
    global _cached_token, _token_expires_at

    if _cached_token and time.time() < _token_expires_at:
        return _cached_token

    cfg = get_config()
    cid = client_id or cfg.keycloak_client_id
    secret = client_secret or cfg.keycloak_client_secret
    url = keycloak_url or cfg.keycloak_url

    if not cid or not secret:
        raise RuntimeError(
            "Keycloak client credentials not configured. "
            "Set ANDLAKE_KEYCLOAK_CLIENT_ID and ANDLAKE_KEYCLOAK_CLIENT_SECRET "
            "environment variables, or pass client_id/client_secret directly."
        )

    token_url = f"{url.rstrip('/')}/protocol/openid-connect/token"
    data = urlencode({
        "grant_type": "client_credentials",
        "client_id": cid,
        "client_secret": secret,
    }).encode()

    req = Request(token_url, data=data, method="POST")
    req.add_header("Content-Type", "application/x-www-form-urlencoded")

    try:
        with urlopen(req, timeout=10) as resp:
            body = json.loads(resp.read())
    except Exception as e:
        raise RuntimeError(f"Failed to fetch Keycloak token: {e}") from e

    _cached_token = body["access_token"]
    # Cache until 60s before expiry
    _token_expires_at = time.time() + body.get("expires_in", 300) - 60
    return _cached_token


def clear_token_cache() -> None:
    """Clear the cached token (useful for testing)."""
    global _cached_token, _token_expires_at
    _cached_token = None
    _token_expires_at = 0.0
