from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Literal

from mcp.server.fastmcp import FastMCP


class MCPServer:
    """
    MCP server for serving a pre-built MCP documentation bundle.
    """

    def __init__(self, mcp_root: Path, name: str) -> None:
        """
        Initialize the MCPServer.

        Args:
            mcp_root: Path to the directory containing pre-built MCP JSON resources.
            name: Name of the MCP server.
        """
        self.mcp_root = mcp_root
        self.app = FastMCP(name)

        self._register_resources()
        self._register_tools()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _read_json(self, path: Path) -> Any:
        """
        Read and parse a JSON file, returning diagnostic errors if missing.

        Args:
            path: Path to the JSON file.

        Returns:
            The parsed JSON data or an error dictionary.
        """
        if not path.exists():
            return {
                "error": "not_found",
                "path": str(path),
            }
        return json.loads(path.read_text(encoding="utf-8"))

    # ------------------------------------------------------------------
    # MCP resources
    # ------------------------------------------------------------------

    def _register_resources(self) -> None:
        """
        Register MCP resources for index, nav, and individual modules.
        """
        @self.app.resource("docs://index")
        def index():
            return self._read_json(self.mcp_root / "index.json")

        @self.app.resource("docs://nav")
        def nav():
            return self._read_json(self.mcp_root / "nav.json")

        @self.app.resource("docs://modules/{module}")
        def module(module: str):
            return self._read_json(
                self.mcp_root / "modules" / f"{module}.json"
            )

    # ------------------------------------------------------------------
    # MCP tools (optional / diagnostic)
    # ------------------------------------------------------------------

    def _register_tools(self) -> None:
        """
        Register high-level MCP tools for diagnostics.
        """
        @self.app.tool()
        def ping() -> str:
            return "pong"

    # ------------------------------------------------------------------
    # Server lifecycle
    # ------------------------------------------------------------------

    def run(self, transport: Literal["stdio", "sse", "streamable-http"] = "streamable-http") -> None:
        """
        Start the MCP server.

        Args:
            transport: MCP transport (default: streamable-http)
        """
        self.app.run(transport=transport)