from __future__ import annotations

import asyncio
from typing import Any

from fastmcp.exceptions import ToolError
from mcp.types import ToolAnnotations

from mcp_eregistrations_keycloak.keycloak_client import (
    KeycloakAdminClient,
    KeycloakClientError,
    translate_error,
)
from mcp_eregistrations_keycloak.server import mcp
from mcp_eregistrations_keycloak.tools.validation import _resolve_realm


@mcp.tool(annotations=ToolAnnotations(readOnlyHint=True))
async def kc_get_token_info(
    instance: str | None = None,
) -> dict[str, Any]:
    """Get metadata about the current session token. Does NOT expose raw tokens.

    Args:
        instance: Instance name, or None for auto-select.

    Returns:
        dict with token_type, expires_in, scope, claims.
    """
    try:
        import jwt as pyjwt
    except ImportError:
        raise ToolError(
            "pyjwt[crypto] is required for JWT tools."
            " Install it with: uv add 'pyjwt[crypto]'"
        )

    from mcp_eregistrations_common.auth import get_or_refresh_token

    token = await get_or_refresh_token(instance) if instance else None
    if not token:
        raise ToolError("Not authenticated. Run kc_auth_login first.")

    try:
        decoded = pyjwt.decode(token, options={"verify_signature": False})
    except Exception as e:
        raise ToolError(f"Failed to decode session token: {e}")

    return {
        "token_type": "Bearer",
        "expires_at": decoded.get("exp"),
        "issued_at": decoded.get("iat"),
        "scope": decoded.get("scope"),
        "claims": {
            "sub": decoded.get("sub"),
            "preferred_username": decoded.get("preferred_username"),
            "email": decoded.get("email"),
            "realm_access": decoded.get("realm_access"),
            "resource_access": decoded.get("resource_access"),
        },
    }


@mcp.tool(annotations=ToolAnnotations(readOnlyHint=True))
async def kc_decode_token(
    token: str,
) -> dict[str, Any]:
    """Decode a JWT token without verification (local decode).

    Args:
        token: JWT token string to decode.

    Returns:
        dict with header, payload.
    """
    try:
        import jwt as pyjwt
    except ImportError:
        raise ToolError(
            "pyjwt[crypto] is required for JWT tools."
            " Install it with: uv add 'pyjwt[crypto]'"
        )

    try:
        header = pyjwt.get_unverified_header(token)
        payload = pyjwt.decode(token, options={"verify_signature": False})
        return {"header": header, "payload": payload}
    except Exception as e:
        raise ToolError(f"Failed to decode token: {e}")


@mcp.tool(annotations=ToolAnnotations(readOnlyHint=True))
async def kc_validate_token(
    token: str,
    realm: str | None = None,
    instance: str | None = None,
) -> dict[str, Any]:
    """Validate a JWT token against Keycloak's JWKS endpoint.

    Args:
        token: JWT token string to validate.
        realm: Realm name, or None for instance default.
        instance: Instance name, or None for auto-select.

    Returns:
        dict with valid, expired, issuer, subject, audience, expires_at, issued_at.
    """
    try:
        import jwt as pyjwt
        from jwt import PyJWKClient
    except ImportError:
        raise ToolError(
            "pyjwt[crypto] is required for JWT tools."
            " Install it with: uv add 'pyjwt[crypto]'"
        )

    try:
        async with KeycloakAdminClient.for_instance(instance) as client:
            r = _resolve_realm(realm, instance, client)
            jwks_url = (
                f"{client.public_base_url}/realms/{r}/protocol/openid-connect/certs"
            )

            try:
                jwks_client = PyJWKClient(jwks_url)
                signing_key = await asyncio.to_thread(
                    jwks_client.get_signing_key_from_jwt, token
                )
                decoded = pyjwt.decode(
                    token,
                    signing_key.key,
                    algorithms=["RS256"],
                    options={"verify_aud": False},
                )
                return {
                    "valid": True,
                    "expired": False,
                    "issuer": decoded.get("iss"),
                    "subject": decoded.get("sub"),
                    "audience": decoded.get("aud"),
                    "expires_at": decoded.get("exp"),
                    "issued_at": decoded.get("iat"),
                }
            except pyjwt.ExpiredSignatureError as e:
                decoded = pyjwt.decode(token, options={"verify_signature": False})
                return {
                    "valid": False,
                    "expired": True,
                    "issuer": decoded.get("iss"),
                    "subject": decoded.get("sub"),
                    "audience": decoded.get("aud"),
                    "expires_at": decoded.get("exp"),
                    "issued_at": decoded.get("iat"),
                    "error": str(e),
                }
            except pyjwt.InvalidTokenError as e:
                return {"valid": False, "expired": False, "error": str(e)}
    except KeycloakClientError as e:
        raise translate_error(e, resource_type="token validation") from e


@mcp.tool(annotations=ToolAnnotations(readOnlyHint=True))
async def kc_introspect_token(
    token: str,
    client_id: str,
    client_secret: str,
    realm: str | None = None,
    instance: str | None = None,
) -> dict[str, Any]:
    """Introspect a token via the Keycloak (RFC 7662).

    Args:
        token: Token to introspect.
        client_id: Client ID for authentication.
        client_secret: Client secret for authentication.
        realm: Realm name, or None for instance default.
        instance: Instance name, or None for auto-select.

    Returns:
        dict with active, scope, client_id, username.
    """
    try:
        async with KeycloakAdminClient.for_instance(instance) as client:
            r = _resolve_realm(realm, instance, client)
            introspect_url = (
                f"{client.public_base_url}/realms/{r}"
                f"/protocol/openid-connect/token/introspect"
            )
            result = await client.post_form(
                introspect_url,
                data={
                    "token": token,
                    "client_id": client_id,
                    "client_secret": client_secret,
                },
            )
            return {
                "active": result.get("active", False),
                "scope": result.get("scope"),
                "client_id": result.get("client_id"),
                "username": result.get("username"),
                "token_type": result.get("token_type"),
                "exp": result.get("exp"),
            }
    except KeycloakClientError as e:
        raise translate_error(e, resource_type="token introspection") from e
