from secrets import compare_digest
from urllib.parse import parse_qs

from starlette.types import ASGIApp, Receive, Scope, Send

PROTECTED_PREFIXES = ("/api", "/mcp")
_UNAUTHORIZED = b'{"detail":"unauthorized"}'


class BearerAuthMiddleware:
    def __init__(self, app: ASGIApp, token: str) -> None:
        self.app = app
        self.token = token

    def _authorized(self, scope: Scope) -> bool:
        headers = dict(scope["headers"])
        bearer = headers.get(b"authorization", b"").decode()
        if bearer.startswith("Bearer ") and compare_digest(bearer[7:], self.token):
            return True
        query = parse_qs(scope["query_string"].decode())
        token = query.get("token", [""])[0]
        return bool(token) and compare_digest(token, self.token)

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return
        if scope["method"] == "OPTIONS" or not scope["path"].startswith(
            PROTECTED_PREFIXES
        ):
            await self.app(scope, receive, send)
            return
        if self._authorized(scope):
            await self.app(scope, receive, send)
            return
        await send(
            {
                "type": "http.response.start",
                "status": 401,
                "headers": [
                    (b"content-type", b"application/json"),
                    (b"www-authenticate", b"Bearer"),
                ],
            }
        )
        await send({"type": "http.response.body", "body": _UNAUTHORIZED})