"""JWTAuth Middleware functionality."""

from typing import Callable

from starlette import status
from starlette.datastructures import Headers
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from jwtauth import JWTAuth
from jwtauth.constants import PROD_ENVIRONMENT
from jwtauth.exceptions import JWTAuthError
from jwtauth.utils import (
    RequestContext,
    get_token_string_from_headers,
    is_path_match_re,
    jwt_auth_enabled_for_env,
    jwt_auth_from_environment,
)


class JWTAuthenticationMiddleware:
    """JWTAuthenticationMiddleware validates/decodes JWT tokens for every request."""

    def __init__(
        self,
        app: ASGIApp,
        environment: str = PROD_ENVIRONMENT,
        enabled: bool | None = None,
        exclude_paths: list[str] | None = None,
        auth: JWTAuth | None = None,
        request_context_func: Callable[[], RequestContext | None] | None = None,
    ) -> None:
        """Create JWTAuthenticationMiddleware.
        app (ASGIApp): reference to the ASGI app
        environment (str): Application environment
        enabled (bool | None): Whether to enable this middleware.
            If nothing is specified, enabled is derived from the environment.
        exclude_paths (list[str] | None): list of paths for which middleware should NOT execute
        auth (JWTAuth | None): Specify a custom JWTAuth object to validate/decode with
        request_context_func (Callable[[], RequestContext | None] | None): Callable to get
            the request context
        """
        self.app = app
        self.environment = environment
        self.enabled = (
            jwt_auth_enabled_for_env(environment) if enabled is None else enabled
        )
        self.exclude_paths = exclude_paths
        self.auth = auth or jwt_auth_from_environment(environment)
        self.request_context_func = request_context_func

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Perform the validation/decoding of JWT token on a request."""
        if (
            scope["type"] != "http"
            or not self.enabled
            or is_path_match_re(scope.get("path", ""), match=self.exclude_paths)
        ):
            await self.app(scope, receive, send)
            return

        try:
            token_string = get_token_string_from_headers(
                Headers(scope=scope),
                request_context_func=self.request_context_func,
            )
            scope["token"] = await self.auth.aget_token(token_string)
        except JWTAuthError as exc:
            response = JSONResponse(
                content={
                    "code": exc.code,
                    "message": exc.message,
                },
                status_code=status.HTTP_401_UNAUTHORIZED,
            )
            await response(scope, receive, send)
        else:
            await self.app(scope, receive, send)