"""Shared authentication for MCP eRegistrations components.

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

from __future__ import annotations

import logging
from typing import Any

import httpx

logger = logging.getLogger(__name__)


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)


# ---------------------------------------------------------------------------
# Token storage (thin wrapper over AuthStore)
# ---------------------------------------------------------------------------


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 persist refresh context to AuthStore.

    Args:
        instance: Instance name (e.g., "jamaica").
        access_token: Access token to cache in memory.
        refresh_token: Refresh token to persist for cross-session auth.
        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 and token_endpoint and client_id:
        from mcp_eregistrations_common.auth_store import store_refresh_context

        try:
            store_refresh_context(
                instance,
                refresh_token=refresh_token,
                token_endpoint=token_endpoint,
                client_id=client_id,
            )
        except Exception:
            pass


# ---------------------------------------------------------------------------
# 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, token_endpoint.

    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()
        result["token_endpoint"] = token_endpoint
        return result


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

    Supports both Keycloak (client_id in body) and CAS (Basic Auth +
    redirect_uri). CAS params are optional — when present, the request
    uses Basic Auth header instead of client_id in body.

    Args:
        token_endpoint: Token endpoint URL.
        client_id: OIDC/OAuth client ID.
        refresh_token: Refresh token.
        client_secret: CAS client secret (enables Basic Auth).
        redirect_uri: CAS redirect URI (required by CAS servers).

    Returns:
        dict with access_token, refresh_token, expires_in.

    Raises:
        AuthenticationError: If refresh fails.
    """
    headers: dict[str, str] = {}
    data: dict[str, str] = {
        "grant_type": "refresh_token",
        "refresh_token": refresh_token,
    }

    if client_secret:
        # CAS: Basic Auth header + redirect_uri (CAS requires redirect_uri
        # on ALL /access_token calls, even refresh_token grants)
        import base64

        creds = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
        headers["Authorization"] = f"Basic {creds}"
        headers["Content-Type"] = "application/x-www-form-urlencoded"
        data["redirect_uri"] = redirect_uri or "http://127.0.0.1:9876/callback"
    else:
        # Keycloak: client_id in body
        data["client_id"] = client_id

    async with httpx.AsyncClient() as client:
        response = await client.post(token_endpoint, headers=headers, data=data)

        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: cache -> refresh -> password grant -> 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

    from mcp_eregistrations_common.auth_store import (
        get_auth_config,
        get_credentials,
        get_refresh_context,
        store_refresh_context,
    )

    # 2. Refresh token
    ctx = get_refresh_context(instance)
    if ctx and ctx.get("refresh_token"):
        try:
            result = await refresh_access_token(
                token_endpoint=ctx["token_endpoint"],
                client_id=ctx["client_id"],
                refresh_token=ctx["refresh_token"],
                client_secret=ctx.get("cas_client_secret"),
                redirect_uri=ctx.get("redirect_uri"),
            )
            store_refresh_context(
                instance,
                refresh_token=result.get("refresh_token", ctx["refresh_token"]),
                token_endpoint=ctx["token_endpoint"],
                client_id=ctx["client_id"],
                cas_client_secret=ctx.get("cas_client_secret"),
                redirect_uri=ctx.get("redirect_uri"),
            )
            access_token: str = result["access_token"]
            _access_tokens[instance] = access_token
            return access_token
        except Exception:
            pass

    # 3. Password grant from stored credentials
    creds = get_credentials(instance)
    auth_cfg = get_auth_config(instance)

    # 3a. Keycloak password grant (OIDC discovery + grant)
    if creds and auth_cfg and auth_cfg.get("keycloak_url"):
        try:
            result = await keycloak_password_grant(
                keycloak_url=auth_cfg["keycloak_url"],
                realm=auth_cfg["keycloak_realm"],
                client_id=auth_cfg.get("client_id", "mcp-eregistrations-bpa"),
                username=creds[0],
                password=creds[1],
            )
            store_refresh_context(
                instance,
                refresh_token=result.get("refresh_token", ""),
                token_endpoint=result["token_endpoint"],
                client_id=auth_cfg.get("client_id", "mcp-eregistrations-bpa"),
            )
            access_token2: str = result["access_token"]
            _access_tokens[instance] = access_token2
            return access_token2
        except Exception:
            pass

    return None
