"""MCP server implementation."""

from collections.abc import Awaitable, Callable
from typing import Any, TypeVar

from fastapi import Body, FastAPI, Header, HTTPException, status
from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, ConfigDict, Field


T = TypeVar("T", bound=BaseModel)


class Tool(BaseModel):
    """MCP tool definition."""

    model_config = ConfigDict(strict=True)

    name: str = Field(description="The name of the tool")
    description: str = Field(description="The description of what the tool does")
    inputSchema: dict[str, Any] = Field(description="The input schema for the tool")
    securitySchemes: list[dict[str, Any]] | None = Field(
        default=None,
        description="Optional security schemes required to invoke this tool",
    )
    strict: bool = Field(
        default=True,
        description="Whether the tool response is strictly validated",
    )
    function: Callable[..., Any] = Field(
        description="The tool function",
        exclude=True,
    )


class OAuth2Configuration(BaseModel):
    """Configuration for describing the OAuth2 / OpenID Connect server."""

    issuer: str = Field(
        description="Issuer identifier for the authorization server",
    )
    authorization_endpoint: str = Field(
        description="OAuth2 / OIDC authorization endpoint URL",
    )
    token_endpoint: str = Field(
        description="OAuth2 / OIDC token endpoint URL",
    )
    registration_endpoint: str | None = Field(
        default=None,
        description="OAuth2 / OIDC registration endpoint URL",
    )
    jwks_uri: str | None = Field(
        default=None,
        description="JWKS endpoint URL if JWT based access tokens are used",
    )
    scopes_supported: list[str] = Field(
        default_factory=list,
        description="All scopes that this authorization server understands",
    )
    token_endpoint_auth_methods_supported: list[str] = Field(
        default_factory=lambda: ["client_secret_basic"],
        description="Supported client authentication methods for the token endpoint",
    )


AuthContext = dict[str, Any]


OAuth2Handler = Callable[[str, list[str]], Awaitable[AuthContext]]


class MCPServer:
    """MCP server implementation."""

    def __init__(
        self,
        name: str,
        description: str | None = None,
        oauth2_configuration: OAuth2Configuration | None = None,
        oauth2_handler: OAuth2Handler | None = None,
    ):
        """Initialize the MCP server."""
        self.name = name
        self.description = (
            description or f"A Model Context Protocol (MCP) tool server for {name}"
        )

        # OAuth2 support
        self._oauth2_configuration = oauth2_configuration
        self._oauth2_handler = oauth2_handler

        # Create FastAPI app
        self.app = FastAPI(
            title=self.name,
            description=self.description,
            version="3.0.0",
        )

        # Initialize MCP server
        self.mcp_server = FastMCP(self.name)

        # Store registered tools
        self._tools: dict[str, Tool] = {}

        # Register endpoints
        self._register_endpoints()

    @staticmethod
    def _sanitize_input_schema(input_schema: dict[str, Any]) -> dict[str, Any]:
        """Remove internal fields like '__auth__' from the public tool schema.

        This ensures that the MCP /tools discovery endpoint does not expose
        the internal authentication context parameter that is injected only
        at invocation time.
        """
        # Shallow copy first so we don't mutate the original schema
        sanitized: dict[str, Any] = {**input_schema}

        properties = sanitized.get("properties")
        if isinstance(properties, dict) and "__auth__" in properties:
            # Remove the internal property
            properties = {k: v for k, v in properties.items() if k != "__auth__"}
            sanitized["properties"] = properties

        required = sanitized.get("required")
        if isinstance(required, list) and "__auth__" in required:
            sanitized["required"] = [name for name in required if name != "__auth__"]

        return sanitized

    def _build_authentication_error(
        self,
        message: str,
        required_scopes: list[str],
    ) -> dict[str, Any]:
        """Build a standardized MCP authentication error response."""
        scope_value = " ".join(required_scopes)

        resource_metadata = ""
        if self._oauth2_configuration is not None and self._oauth2_configuration.issuer:
            issuer = self._oauth2_configuration.issuer.rstrip("/")
            resource_metadata = f"{issuer}/.well-known/oauth-protected-resource"

        www_authenticate = (
            'Bearer '
            f'resource_metadata="{resource_metadata}", '
            f'scope="{scope_value}", '
            'error="invalid_token", '
            'error_description="User must sign in first."'
        )

        return {
            "isError": True,
            "content": [
                {
                    "type": "text",
                    "text": message,
                }
            ],
            "_meta": {
                "mcp/www_authenticate": www_authenticate,
            },
        }

    def _register_endpoints(self) -> None:
        """Register FastAPI endpoints."""

        @self.app.get("/tools", response_model=list[Tool])
        async def get_tools() -> list[Tool]:
            """Get all available MCP tools."""
            return list(self._tools.values())

        @self.app.post("/tools/{tool_name}/invoke")
        async def invoke_tool(
            tool_name: str,
            arguments: dict[str, Any] = Body(...),
            authorization: str | None = Header(default=None),
        ):
            """Invoke a specific MCP tool."""
            tool = self._tools.get(tool_name)
            if tool is None:
                raise HTTPException(
                    status_code=status.HTTP_404_NOT_FOUND,
                    detail=f"Unknown tool '{tool_name}'",
                )

            # Enforce OAuth2 if the tool defines security schemes
            required_scopes: list[str] = []
            auth_context: AuthContext | None = None
            if tool.securitySchemes:
                for scheme in tool.securitySchemes:
                    if scheme.get("type") == "oauth2":
                        required_scopes.extend(scheme.get("scopes", []))

                if required_scopes:
                    if self._oauth2_handler is None:
                        raise HTTPException(
                            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                            detail=(
                                "Tool requires OAuth2 but no oauth2_handler "
                                "is configured for MCPServer"
                            ),
                        )
                    if not authorization or not authorization.startswith("Bearer "):
                        return self._build_authentication_error(
                            "Authentication required to invoke this tool.",
                            required_scopes,
                        )

                    token = authorization.removeprefix("Bearer ").strip()
                    # Delegate token validation and scope checks to the configured handler.
                    # The handler can return an authentication context (e.g. subject, scopes),
                    # which will be injected into the tool arguments.
                    try:
                        auth_context = await self._oauth2_handler(
                            token,
                            required_scopes,
                        )
                    except HTTPException as exc:
                        if exc.status_code in {
                            status.HTTP_401_UNAUTHORIZED,
                            status.HTTP_403_FORBIDDEN,
                        }:
                            return self._build_authentication_error(
                                str(exc.detail)
                                if isinstance(exc.detail, str)
                                else "Authentication required to invoke this tool.",
                                required_scopes,
                            )
                        raise

            if auth_context is not None:
                # Make the authentication context available to the tool implementation.
                # Tools that benötigen Auth-Informationen können ein entsprechendes
                # Feld im Input-Modell definieren (z. B. "__auth__").
                arguments["__auth__"] = auth_context

            try:
                result = await self.mcp_server.call_tool(
                    tool_name,
                    arguments=arguments,
                )
                if isinstance(result, list) and len(result) == 1:
                    first_item = result[0]
                    if hasattr(first_item, "type") and first_item.type == "text":
                        return first_item.text
                elif isinstance(result, dict) and result.get("type") == "text":
                    return result["text"]
                return result
            except HTTPException:
                # Re-raise HTTP exceptions untouched
                raise
            except Exception as exc:  # pragma: no cover - safety net
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=f"Error invoking tool {tool_name}: {exc}",
                ) from exc

        # OAuth2 / OpenID configuration discovery
        if self._oauth2_configuration is not None:

            @self.app.get("/.well-known/openid-configuration")
            async def openid_configuration() -> dict[str, Any]:
                """Expose OAuth2 / OpenID Connect discovery document."""
                return self._oauth2_configuration.model_dump()

            @self.app.get("/.well-known/oauth-authorization-server")
            async def oauth_authorization_server() -> dict[str, Any]:
                """Expose OAuth2 authorization server metadata."""
                return self._oauth2_configuration.model_dump()

    async def add_tool(
        self,
        tool: Callable[..., Any],
        security_schemes: list[dict[str, Any]] | None = None,
    ) -> None:
        """Add a tool to the MCP server."""
        # Add tool to MCP server
        self.mcp_server.add_tool(tool)

        # Get tool metadata
        tools = await self.mcp_server.list_tools()
        tool_metadata = tools[-1]

        # Use the raw inputSchema from FastMCP - it will be processed later by convert_to_strict_schema
        # when the tool is converted to OpenAI function format
        raw_input_schema = tool_metadata.inputSchema
        input_schema = self._sanitize_input_schema(raw_input_schema)

        # Create and store Tool instance
        self._tools[tool_metadata.name] = Tool(
            name=tool_metadata.name,
            description=tool_metadata.description,
            inputSchema=input_schema,
            securitySchemes=security_schemes,
            strict=True,
            function=tool,
        )