"""Auth0 authentication helpers for role-based access control.

This module provides utilities to check user roles from Auth0 JWT tokens.
Works with both the official MCP SDK and FastMCP.

Example:
    from point_topic_mcp.auth import require_admin_role

    @mcp.tool()
    async def admin_only_tool(token_claims: dict | None = None):
        require_admin_role(token_claims)  # Raises PermissionError if not admin
        return {"message": "Admin access granted"}
"""

import logging
from typing import Any

import httpx
from jose import jwt, JWTError

logger = logging.getLogger(__name__)

# Thread-local storage for current token context
# This allows token to be passed without changing function signatures everywhere
_current_token_claims: dict[str, Any] = {}


def set_current_token_claims(claims: dict | None):
    """Set the current token claims for this request context.

    This is used internally by the server to make token claims
    available to tools without changing their signatures.
    """
    global _current_token_claims
    _current_token_claims = claims or {}


def get_current_token_claims() -> dict | None:
    """Get the current token claims for this request context.

    Returns:
        Token claims dict if available, None otherwise
    """
    return _current_token_claims if _current_token_claims else None


def get_user_roles_from_claims(claims: dict | None = None) -> list[str]:
    """Get roles from JWT token claims.

    Extracts roles from Auth0 JWT token claims. Auth0 can store roles
    in different claim locations depending on the token configuration.

    Args:
        claims: Token claims dict. If None, uses current context claims.

    Returns:
        List of role/permission strings from the token
        Empty list if no token or no roles found
    """
    if claims is None:
        claims = get_current_token_claims()

    if not claims:
        return []

    # Auth0 stores roles in different places depending on configuration
    # Try common locations in order of preference

    # 1. Direct 'roles' claim (configured in Auth0 Actions/rules)
    roles = claims.get("roles", [])
    if roles:
        return roles if isinstance(roles, list) else [roles]

    # 2. 'permissions' claim (from Auth0 API Authorization)
    permissions = claims.get("permissions", [])
    if permissions:
        return permissions if isinstance(permissions, list) else [permissions]

    # 3. 'scope' claim (OAuth2 scope, space-separated)
    scope = claims.get("scope", "")
    if scope:
        return scope.split() if isinstance(scope, str) else [scope]

    return []


# Backwards compatibility alias
get_user_roles = get_user_roles_from_claims


def has_role(role_name: str, claims: dict | None = None) -> bool:
    """Check if user has a specific role.

    Args:
        role_name: The role to check for
        claims: Token claims dict. If None, uses current context claims.

    Returns:
        True if user has the role, False otherwise
    """
    roles = get_user_roles_from_claims(claims)
    return role_name in roles


def has_any_role(role_names: list[str], claims: dict | None = None) -> bool:
    """Check if user has any of the specified roles.

    Args:
        role_names: List of roles to check for
        claims: Token claims dict. If None, uses current context claims.

    Returns:
        True if user has at least one of the roles
    """
    roles = get_user_roles_from_claims(claims)
    return any(role in roles for role in role_names)


def require_admin_role(claims: dict | None = None):
    """Require user to have Point Topic Admin role.

    Checks for admin roles. Raises PermissionError if user
    doesn't have appropriate admin access.

    Args:
        claims: Token claims dict. If None, uses current context claims.

    Raises:
        PermissionError: If user doesn't have admin role
    """
    roles = get_user_roles_from_claims(claims)

    # Admin role names to check for
    admin_roles = [
        "Point Topic Admin",
        "ptAdmin",
        "admin",
        "user:upc-query-agent",  # From Auth0 permissions
        "admin:upc-query-agent",  # From Auth0 permissions
    ]

    if not any(role in roles for role in admin_roles):
        raise PermissionError(
            f"This tool requires admin role. "
            f"Your roles: {roles if roles else 'None'}. "
            f"Required: one of {admin_roles}"
        )


def require_role(role_name: str, claims: dict | None = None):
    """Require user to have a specific role.

    Generic role checker. Raises PermissionError if user
    doesn't have the specified role.

    Args:
        role_name: The required role name
        claims: Token claims dict. If None, uses current context claims.

    Raises:
        PermissionError: If user doesn't have the role
    """
    roles = get_user_roles_from_claims(claims)

    if role_name not in roles:
        raise PermissionError(
            f"This tool requires '{role_name}' role. "
            f"Your roles: {roles if roles else 'None'}"
        )


def get_user_info(claims: dict | None = None) -> dict:
    """Get user information from the JWT token.

    Args:
        claims: Token claims dict. If None, uses current context claims.

    Returns:
        Dictionary with user claims (sub, email, name, etc.)
        Empty dict if no token available
    """
    if claims is None:
        claims = get_current_token_claims()

    if not claims:
        return {}

    return {
        "user_id": claims.get("sub"),
        "email": claims.get("email"),
        "name": claims.get("name"),
        "roles": get_user_roles_from_claims(claims),
        "permissions": claims.get("permissions", []),
    }


class Auth0TokenVerifier:
    """JWT token verifier using Auth0 JWKS endpoint.

    This class validates JWT tokens issued by Auth0 using the
    JSON Web Key Set (JWKS) endpoint for signature verification.
    """

    def __init__(
        self,
        auth0_domain: str,
        audience: str,
        jwks_url: str | None = None,
    ):
        """Initialize the token verifier.

        Args:
            auth0_domain: Auth0 tenant domain (e.g., 'point-topic.eu.auth0.com')
            audience: Expected audience/resource URL (e.g., 'https://point-topic-mcp.fastmcp.app/')
            jwks_url: Optional JWKS URL. Defaults to https://{domain}/.well-known/jwks.json
        """
        self.auth0_domain = auth0_domain
        self.audience = audience
        self.jwks_url = jwks_url or f"https://{auth0_domain}/.well-known/jwks.json"
        self._jwks: dict | None = None

    async def _get_jwks(self) -> dict | None:
        """Fetch JWKS from Auth0 if not already cached."""
        if self._jwks is None:
            async with httpx.AsyncClient() as client:
                response = await client.get(self.jwks_url)
                response.raise_for_status()
                self._jwks = response.json()
        return self._jwks

    async def verify_token(self, token: str) -> dict | None:
        """Verify and decode a JWT token.

        Args:
            token: The JWT token string to verify

        Returns:
            Decoded token claims dict if valid, None otherwise
        """
        try:
            jwks = await self._get_jwks()
            if jwks is None:
                logger.warning("Failed to fetch JWKS")
                return None
            signing_key = self._get_signing_key(token, jwks)

            if not signing_key:
                logger.warning("No signing key found for token")
                return None

            # Decode and verify the token
            payload = jwt.decode(
                token,
                signing_key,
                audience=self.audience,
                issuer=f"https://{self.auth0_domain}/",
            )

            return payload

        except jwt.ExpiredSignatureError:  # type: ignore
            logger.warning("Token has expired")
        except jwt.JWTClaimsError as e:  # type: ignore
            logger.warning(f"Token claims error: {e}")
        except JWTError as e:
            logger.warning(f"JWT validation error: {e}")
        except Exception as e:
            logger.error(f"Unexpected error verifying token: {e}")

        return None

    def _get_signing_key(self, token: str, jwks: dict) -> dict | None:
        """Extract the signing key from JWKS matching the token's kid."""
        try:
            headers = jwt.get_unverified_header(token)
            kid = headers.get("kid")
            if not kid:
                return None

            for key in jwks.get("keys", []):
                if key.get("kid") == kid:
                    return key
        except JWTError as e:
            logger.warning(f"Failed to extract token header: {e}")

        return None

    async def get_token_claims(self, token: str) -> dict | None:
        """Get claims from a token (wrapper around verify_token).

        Args:
            token: The JWT token string

        Returns:
            Token claims dict if valid, None otherwise
        """
        return await self.verify_token(token)


__all__ = [
    # Token context management
    "set_current_token_claims",
    "get_current_token_claims",
    # Role checking functions
    "get_user_roles",
    "get_user_roles_from_claims",
    "has_role",
    "has_any_role",
    "require_admin_role",
    "require_role",
    "get_user_info",
    # Token verification
    "Auth0TokenVerifier",
]
