"""Restrict access to specific users or groups from the identity provider."""

from __future__ import annotations

import logging
from typing import Any

from fastmcp.exceptions import ToolError
from fastmcp.server.dependencies import get_access_token
from fastmcp.server.middleware import Middleware, MiddlewareContext

logger = logging.getLogger("bookstack-mcp.authz")

# Claims that different providers use for the human-readable username.
# GitHub -> login;  Keycloak/Authentik/Authelia/Zitadel -> preferred_username.
USERNAME_CLAIMS = ("login", "preferred_username", "username", "email", "sub")


def _dig(claims: dict[str, Any], path: str) -> Any:
    """Look up a possibly nested claim, e.g. 'resource_access.mcp.roles'."""
    value: Any = claims
    for part in path.split("."):
        if not isinstance(value, dict) or part not in value:
            return None
        value = value[part]
    return value


class IdentityAllowlistMiddleware(Middleware):
    """Reject requests from identities outside the configured allowlist.

    Two independent checks, both optional:
      * allowed_users   -- username must appear in this list
      * required_groups -- at least one group/role claim must match

    Static machine tokens carry no identity claims. Those are already trusted
    because the operator minted them by hand, so they pass through untouched.
    """

    def __init__(
        self,
        allowed_users: list[str] | None = None,
        required_groups: list[str] | None = None,
        groups_claim: str = "groups",
    ) -> None:
        self.allowed_users = {u.lower() for u in (allowed_users or [])}
        self.required_groups = {g.lower() for g in (required_groups or [])}
        self.groups_claim = groups_claim

    @staticmethod
    def _username(claims: dict[str, Any]) -> str | None:
        for claim in USERNAME_CLAIMS:
            value = claims.get(claim)
            if isinstance(value, str) and value:
                return value
        return None

    def _groups(self, claims: dict[str, Any]) -> set[str]:
        raw = _dig(claims, self.groups_claim)
        if isinstance(raw, str):
            raw = [raw]
        if not isinstance(raw, list):
            return set()
        return {str(g).lower() for g in raw}

    def _check(self) -> None:
        token = get_access_token()
        if token is None:
            return

        claims = token.claims or {}
        username = self._username(claims)
        if username is None:
            return  # static token, no identity to check

        if self.allowed_users and username.lower() not in self.allowed_users:
            logger.warning("Denied BookStack MCP access to user %r (not in allowlist)", username)
            raise ToolError(f"User '{username}' is not authorised to use this BookStack server.")

        if self.required_groups:
            groups = self._groups(claims)
            if not (groups & self.required_groups):
                logger.warning(
                    "Denied BookStack MCP access to user %r (groups %s, need one of %s)",
                    username, sorted(groups), sorted(self.required_groups),
                )
                raise ToolError(
                    f"User '{username}' is not in a group permitted to use this BookStack server."
                )

    async def on_call_tool(self, context: MiddlewareContext, call_next):
        self._check()
        return await call_next(context)

    async def on_list_tools(self, context: MiddlewareContext, call_next):
        self._check()
        return await call_next(context)