from asgiref.sync import sync_to_async
from rest_framework.exceptions import APIException
from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.exceptions import TokenError

from mcp.server.auth.provider import AccessToken
from mcp.server.auth.middleware.bearer_auth import RequireAuthMiddleware
from starlette.types import ASGIApp, Receive, Scope, Send


def verify_factory_jwt(token: str) -> AccessToken | None:
    """验证 Factory access token，并生成 MCP 的逐请求身份信息。"""
    authentication = JWTAuthentication()
    try:
        validated_token = authentication.get_validated_token(token)
        user = authentication.get_user(validated_token)
    except (APIException, TokenError):
        return None

    return AccessToken(
        token=token,
        client_id="factory-mcp",
        scopes=["factory:user"],
        expires_at=validated_token.get("exp"),
        subject=str(user.pk),
        claims={
            "factory_user": {
                "id": str(user.pk),
                "username": user.get_username(),
                "name": user.name,
                "is_superuser": user.is_superuser,
            },
        },
    )


class FactoryJWTVerifier:
    """让 MCP SDK 复用 Factory SimpleJWT 的验证规则。"""

    async def verify_token(self, token: str) -> AccessToken | None:
        return await sync_to_async(
            verify_factory_jwt,
            thread_sensitive=True,
        )(token)


class RequireFactoryJWTMiddleware:
    """仅保护 HTTP 请求，并把 ASGI lifespan 原样交给 MCP SDK。"""

    def __init__(self, app: ASGIApp):
        self.app = app
        self.protected_app = RequireAuthMiddleware(
            app,
            required_scopes=[],
        )

    async def __call__(
        self,
        scope: Scope,
        receive: Receive,
        send: Send,
    ) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return
        await self.protected_app(scope, receive, send)