"""
Security module for FastAPI to handle JWT token verification.
"""

import re
from typing import Awaitable, Callable, Iterable, Optional

from cachetools import LRUCache
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp

from monday_com_orca_backend import env_vars
from monday_com_orca_backend.api.auth import utils
from monday_com_orca_backend.enums import HttpHeader, HttpMethod, JwtAlgorithm
from monday_com_orca_backend.utils.custom_types import FrozenDict

from .exceptions import UnauthorizedException
from .typings import ErrorHandler

JWT_ALGORITHM: JwtAlgorithm = JwtAlgorithm.HS256
BEARER_TOKEN_REGEX: re.Pattern = re.compile(r"^Bearer (?P<token>.+)$", re.IGNORECASE)
VALID_TOKEN_CACHE: LRUCache = LRUCache(maxsize=4)


class JWTMiddleware(BaseHTTPMiddleware):
    """FastAPI Middleware to validate JWT tokens on incoming requests,
    except for specified paths or path prefixes.
    """

    def __init__(
        self,
        app: ASGIApp,
        excluded_paths: Optional[Iterable[str]] = None,
        excluded_prefixes: Optional[Iterable[str]] = None,
        custom_error_handler: Optional[ErrorHandler] = None,
        ignore_trailing_slash: bool = False,
    ) -> None:
        """
        Initialize the JWTMiddleware.

        Args:
            app: FastAPI application instance.
            excluded_paths: A list of exact paths to exclude from JWT validation.
              E.g. ["/public", "/static/"] will exclude the exact paths "/public"
              and "/static/" from JWT validation, but not their subpaths like
              "/public/page" or "/static/css/".
            excluded_prefixes: A list of path prefixes to exclude from JWT validation.
              E.g. ["/public", "/static/"] will exclude all paths starting with
              "/public" or "/static/", including their subpaths like
              "/public/page" or "/static/css/".
            ignore_trailing_slash: If True, ignore trailing slashes in paths.
              E.g. "/public/" and "/public" will be treated as the same path.
        """
        super().__init__(app)
        self._excluded_paths: set[str] = set()
        self._excluded_prefixes: set[str] = set()
        self._original_excluded_paths: list[str] = []
        self._original_excluded_prefixes: list[str] = []
        self._ignore_trailing_slash: bool = ignore_trailing_slash
        self.custom_error_handler: Optional[ErrorHandler] = custom_error_handler
        self._set_excluded_paths_and_prefixes(excluded_paths, excluded_prefixes)

    def _normalize_path(self, path: str) -> str:
        """Normalize a path by removing trailing slash if ignore_trailing_slash
        is set.
        """
        if self.ignore_trailing_slash and path != "/":
            return path.rstrip("/")
        return path

    def _set_excluded_paths_and_prefixes(
        self,
        excluded_paths: Optional[Iterable[str]],
        excluded_prefixes: Optional[Iterable[str]],
    ):
        self._excluded_paths = {self._normalize_path(p) for p in (excluded_paths or [])}
        self._excluded_prefixes = {
            self._normalize_path(p) for p in (excluded_prefixes or [])
        }
        # Store originals for property setters
        self._original_excluded_paths = list(excluded_paths) if excluded_paths else []
        self._original_excluded_prefixes = (
            list(excluded_prefixes) if excluded_prefixes else []
        )

    @property
    def ignore_trailing_slash(self) -> bool:
        return self._ignore_trailing_slash

    @ignore_trailing_slash.setter
    def ignore_trailing_slash(self, value: bool):
        self._ignore_trailing_slash = value
        # Renormalize excluded paths and prefixes if the flag changes
        self._set_excluded_paths_and_prefixes(
            self.excluded_paths, self.excluded_prefixes
        )

    @property
    def excluded_paths(self) -> set[str]:
        return self._excluded_paths

    @excluded_paths.setter
    def excluded_paths(self, value):
        self._set_excluded_paths_and_prefixes(value, self._original_excluded_prefixes)

    @property
    def excluded_prefixes(self) -> set[str]:
        return self._excluded_prefixes

    @excluded_prefixes.setter
    def excluded_prefixes(self, value):
        self._set_excluded_paths_and_prefixes(self._original_excluded_paths, value)

    async def dispatch(
        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
    ) -> Response:
        path = self._normalize_path(request.url.path)
        if path in self.excluded_paths:
            return await call_next(request)
        if any(path.startswith(prefix) for prefix in self.excluded_prefixes):
            return await call_next(request)
        try:
            await verify_token(request)
        except UnauthorizedException as exc:
            if self.custom_error_handler:
                return await self.custom_error_handler(request, exc)
            raise exc from None
        return await call_next(request)


async def verify_token(request: Request) -> None:
    """Verify the JWT token in the Authorization header.

    This function can be used as a FastAPI dependency to protect specific routes.
    """
    if request.method == HttpMethod.OPTIONS:
        return None
    auth_header = request.headers.get(HttpHeader.AUTHORIZATION, "")
    try:
        cache_key = auth_header.strip()
    except AttributeError:
        # Token is not a string, but should be
        raise UnauthorizedException("Authorization header is malformed") from None
    cached_result = VALID_TOKEN_CACHE.get(cache_key)  # Use cache for performance
    if cached_result is not None:
        match, payload = cached_result
    else:
        match = BEARER_TOKEN_REGEX.match(auth_header)
        token = match.group("token") if match else None
        if not token:
            raise UnauthorizedException(
                "Authorization header is missing or malformed"
            ) from None
        payload = utils.validate_jwt_token(
            token, env_vars.MONDAY_APP_CLIENT_SECRET, algorithm=JWT_ALGORITHM
        )
        if payload is None:
            raise UnauthorizedException("Invalid Signing Secret Token") from None

        # Make payload read-only if not already
        payload = FrozenDict(payload) if isinstance(payload, dict) else payload
        VALID_TOKEN_CACHE[cache_key] = (match, payload)
    request.state.user_payload = payload
    return None