"""Shared authentication for MCP eRegistrations components.

Provides:
- Keycloak password grant (used by all components)
- Keyring-based token sharing (access + refresh tokens)
- Cross-server token resolution (login once, use everywhere)
"""

from __future__ import annotations

import json
import logging
from typing import Any

import httpx

logger = logging.getLogger(__name__)

KEYRING_SERVICE = "mcp-eregistrations"


class AuthenticationError(Exception):
    """Authentication failed."""


# ---------------------------------------------------------------------------
# In-memory token cache (per-process)
# ---------------------------------------------------------------------------

_access_tokens: dict[str, str] = {}


def get_cached_token(instance: str) -> str | None:
    """Get access token from in-memory cache."""
    return _access_tokens.get(instance)


# ---------------------------------------------------------------------------
# Keyring operations
# ---------------------------------------------------------------------------


def _get_keyring() -> Any:
    """Lazy-import keyring to avoid hard dependency at module load."""
    try:
        import keyring

        return keyring
    except ImportError:
        return None


def store_shared_tokens(
    instance: str,
    *,
    access_token: str,
    refresh_token: str | None = None,
    token_endpoint: str | None = None,
    client_id: str | None = None,
) -> None:
    """Store tokens in memory and optionally in keyring.

    Args:
        instance: Instance name (e.g., "jamaica").
        access_token: Access token to cache in memory.
        refresh_token: Refresh token to store in keyring.
        token_endpoint: Token endpoint URL (stored with refresh token).
        client_id: Client ID (stored with refresh token).
    """
    _access_tokens[instance] = access_token

    if refresh_token:
        kr = _get_keyring()
        if kr is None:
            logger.debug("keyring not available; refresh token not persisted")
            return
        try:
            payload = json.dumps(
                {
                    "refresh_token": refresh_token,
                    "token_endpoint": token_endpoint,
                    "client_id": client_id,
                }
            )
            kr.set_password(KEYRING_SERVICE, f"{instance}:refresh", payload)
        except Exception as e:
            logger.warning("Cannot store refresh token in keyring: %s", e)


def get_shared_refresh_context(instance: str) -> dict[str, str] | None:
    """Get refresh context from keyring (cross-server).

    Args:
        instance: Instance name.

    Returns:
        Dict with refresh_token, token_endpoint, client_id, or None.
    """
    kr = _get_keyring()
    if kr is None:
        return None
    try:
        raw: str | None = kr.get_password(KEYRING_SERVICE, f"{instance}:refresh")
        if raw:
            result: dict[str, str] = json.loads(raw)
            return result
        return None
    except Exception as e:
        logger.warning("Cannot read refresh token from keyring: %s", e)
        return None


# ---------------------------------------------------------------------------
# Keycloak operations
# ---------------------------------------------------------------------------


async def keycloak_password_grant(
    *,
    keycloak_url: str,
    realm: str,
    client_id: str,
    username: str,
    password: str,
) -> dict[str, Any]:
    """Authenticate via Keycloak password grant.

    Args:
        keycloak_url: Keycloak base URL.
        realm: Keycloak realm.
        client_id: OIDC client ID.
        username: Username.
        password: Password.

    Returns:
        dict with access_token, refresh_token, expires_in.

    Raises:
        AuthenticationError: If discovery or auth fails.
    """
    well_known = (
        f"{keycloak_url.rstrip('/')}/realms/{realm}/.well-known/openid-configuration"
    )

    async with httpx.AsyncClient() as client:
        # Discover token endpoint
        discovery = await client.get(well_known)
        if discovery.status_code != 200:
            raise AuthenticationError(
                f"OIDC discovery failed at {well_known} "
                f"(HTTP {discovery.status_code}). "
                "Check keycloak_url and realm."
            )
        token_endpoint = discovery.json().get("token_endpoint")
        if not token_endpoint:
            raise AuthenticationError("OIDC discovery response missing token_endpoint.")

        # Password grant
        token_response = await client.post(
            token_endpoint,
            data={
                "grant_type": "password",
                "client_id": client_id,
                "username": username,
                "password": password,
            },
        )

        if token_response.status_code != 200:
            error = "Unknown error"
            try:
                error = token_response.json().get("error_description", error)
            except Exception:
                error = token_response.text
            raise AuthenticationError(
                f"Authentication failed for '{username}': {error}"
            )

        result: dict[str, Any] = token_response.json()
        return result


async def refresh_access_token(
    *,
    token_endpoint: str,
    client_id: str,
    refresh_token: str,
) -> dict[str, Any]:
    """Exchange refresh token for new access token.

    Args:
        token_endpoint: Keycloak token endpoint URL.
        client_id: OIDC client ID.
        refresh_token: Refresh token.

    Returns:
        dict with access_token, refresh_token, expires_in.

    Raises:
        AuthenticationError: If refresh fails.
    """
    async with httpx.AsyncClient() as client:
        response = await client.post(
            token_endpoint,
            data={
                "grant_type": "refresh_token",
                "client_id": client_id,
                "refresh_token": refresh_token,
            },
        )

        if response.status_code != 200:
            error = "Unknown error"
            try:
                error = response.json().get("error_description", error)
            except Exception:
                error = response.text
            raise AuthenticationError(f"Token refresh failed: {error}")

        result: dict[str, Any] = response.json()
        return result


async def get_or_refresh_token(instance: str) -> str | None:
    """Get a valid access token: in-memory cache -> keyring refresh -> None.

    Args:
        instance: Instance name.

    Returns:
        Access token string, or None if no token available.
    """
    # 1. In-memory cache
    token = _access_tokens.get(instance)
    if token:
        return token

    # 2. Keyring refresh
    ctx = get_shared_refresh_context(instance)
    if ctx:
        try:
            result = await refresh_access_token(
                token_endpoint=ctx["token_endpoint"],
                client_id=ctx["client_id"],
                refresh_token=ctx["refresh_token"],
            )
            access_token: str = result["access_token"]
            # Update shared tokens (cache + keyring)
            store_shared_tokens(
                instance,
                access_token=access_token,
                refresh_token=result.get("refresh_token"),
                token_endpoint=ctx["token_endpoint"],
                client_id=ctx["client_id"],
            )
            return access_token
        except (AuthenticationError, KeyError, Exception) as e:
            logger.info("Keyring token refresh failed for '%s': %s", instance, e)

    return None
