#!/usr/bin/env python3
"""
Example showing how to create a custom BearerAuthProvider that supports unsigned JWT tokens.

This demonstrates the configuration needed for FastMCP to work with tokens generated
by oauth_token.py which uses alg: "none" and issuer: "local".
"""

import time
from typing import Any

from authlib.jose import JsonWebToken
from authlib.jose.errors import JoseError
from fastmcp.server.auth.auth import (
    ClientRegistrationOptions,
    OAuthProvider,
    RevocationOptions,
)
from fastmcp.utilities.logging import get_logger
from mcp.server.auth.provider import AccessToken
from mcp.shared.auth import OAuthClientInformationFull
from pydantic import AnyHttpUrl, ValidationError


class UnsignedBearerAuthProvider(OAuthProvider):
    """
    Custom Bearer Token validator that supports unsigned JWT tokens (alg: "none").

    This provider is specifically designed to work with tokens generated by oauth_token.py
    which creates unsigned tokens with issuer "local".
    """

    def __init__(
        self,
        issuer: str | None = None,
        audience: str | list[str] | None = None,
        required_scopes: list[str] | None = None,
    ):
        """
        Initialize the provider for unsigned tokens.

        Args:
            issuer: Expected issuer claim (e.g., "local")
            audience: Expected audience claim - can be a string or list of strings
            required_scopes: List of required scopes for access (optional)
        """
        # Handle issuer URL validation for parent class
        try:
            issuer_url = AnyHttpUrl(issuer) if issuer else "https://fastmcp.example.com"
        except ValidationError:
            # Issuer is not a valid URL, use default for parent class
            issuer_url = "https://fastmcp.example.com"

        super().__init__(
            issuer_url=issuer_url,
            client_registration_options=ClientRegistrationOptions(enabled=False),
            revocation_options=RevocationOptions(enabled=False),
            required_scopes=required_scopes,
        )

        self.issuer = issuer
        self.audience = audience
        self.jwt = JsonWebToken(["none"])  # Support only "none" algorithm
        self.logger = get_logger(__name__)

    async def load_access_token(self, token: str) -> AccessToken | None:
        """
        Validates the provided unsigned JWT bearer token.

        Args:
            token: The JWT token string to validate

        Returns:
            AccessToken object if valid, None if invalid or expired
        """
        try:
            # Decode the unsigned JWT token (no key needed for alg: "none")
            claims = self.jwt.decode(token, key="")

            # Extract client ID early for logging
            client_id = claims.get("client_id") or claims.get("sub") or "unknown"

            # Validate expiration
            exp = claims.get("exp")
            if exp and exp < time.time():
                self.logger.debug(
                    "Token validation failed: expired token for client %s", client_id
                )
                self.logger.info("Bearer token rejected for client %s", client_id)
                return None

            # Validate issuer if configured
            if self.issuer:
                if claims.get("iss") != self.issuer:
                    self.logger.debug(
                        "Token validation failed: issuer mismatch for client %s (expected %s, got %s)",
                        client_id,
                        self.issuer,
                        claims.get("iss"),
                    )
                    self.logger.info("Bearer token rejected for client %s", client_id)
                    return None

            # Validate audience if configured
            if self.audience:
                aud = claims.get("aud")
                audience_valid = False

                if isinstance(self.audience, list):
                    if isinstance(aud, list):
                        audience_valid = any(
                            expected in aud for expected in self.audience
                        )
                    else:
                        audience_valid = aud in self.audience
                else:
                    if isinstance(aud, list):
                        audience_valid = self.audience in aud
                    else:
                        audience_valid = aud == self.audience

                if not audience_valid:
                    self.logger.debug(
                        "Token validation failed: audience mismatch for client %s (expected %s, got %s)",
                        client_id,
                        self.audience,
                        aud,
                    )
                    self.logger.info("Bearer token rejected for client %s", client_id)
                    return None

            # Extract scopes
            scopes = self._extract_scopes(claims)

            self.logger.info("Bearer token accepted for client %s", client_id)
            return AccessToken(
                token=token,
                client_id=str(client_id),
                scopes=scopes,
                expires_at=int(exp) if exp else None,
            )

        except JoseError as e:
            self.logger.debug("Token validation failed: JWT format invalid: %s", str(e))
            return None
        except Exception as e:
            self.logger.debug("Token validation failed: %s", str(e))
            return None

    def _extract_scopes(self, claims: dict[str, Any]) -> list[str]:
        """Extract scopes from JWT claims."""
        for claim in ["scope", "scp"]:
            if claim in claims:
                if isinstance(claims[claim], str):
                    return claims[claim].split()
                elif isinstance(claims[claim], list):
                    return claims[claim]
        return []

    async def verify_token(self, token: str) -> AccessToken | None:
        """Verify a bearer token and return access info if valid."""
        return await self.load_access_token(token)

    # --- Unused OAuth server methods (required by interface) ---
    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        raise NotImplementedError("Client management not supported")

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        raise NotImplementedError("Client registration not supported")

    # ... other OAuth methods would be implemented similarly with NotImplementedError


# Example usage:
def create_unsigned_bearer_provider():
    """Create a bearer auth provider for unsigned tokens with issuer 'local'."""
    return UnsignedBearerAuthProvider(
        issuer="local",
        audience="gx-mcp-server",  # Match the audience in oauth_token.py
    )


if __name__ == "__main__":
    # Example test
    import base64
    import json

    # Create test token (like oauth_token.py does)
    header = {"typ": "JWT", "alg": "none"}
    payload = {
        "iss": "local",
        "sub": "demo-client",
        "aud": "gx-mcp-server",
        "iat": int(time.time()),
        "exp": int(time.time()) + 3600,
        "client_id": "demo-client",
    }

    header_b64 = (
        base64.urlsafe_b64encode(json.dumps(header).encode()).decode().rstrip("=")
    )
    payload_b64 = (
        base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
    )
    test_token = f"{header_b64}.{payload_b64}."

    # Test provider
    import asyncio

    async def test():
        provider = create_unsigned_bearer_provider()
        result = await provider.load_access_token(test_token)
        if result:
            print(f"✅ Token validated successfully for client: {result.client_id}")
        else:
            print("❌ Token validation failed")

    asyncio.run(test())
