"""
MCP Router module.

Routes requests within the MCP subsystem, handling tool registration and execution.
"""

from __future__ import annotations

from typing import Optional, Any

from fastapi import APIRouter, Request, Response
from mcp.server.fastmcp import FastMCP
from mcp.server.sse import SseServerTransport


class MCPAsgiResponse(Response):
    """
    Delegates the raw ASGI response cycle to the MCP underlying transport.
    Prevents FastAPI from trying to send a second response.
    """

    def __init__(self, transport):
        # We don't call super().__init__ because we override __call__ entirely
        self.transport = transport

    async def __call__(self, scope, receive, send):
        await self.transport.handle_post_message(scope, receive, send)


class BaseMCPRouter:
    """
    A FastAPI-compatible router that wraps a FastMCP instance.
    """

    def __init__(self, name: str, version: str = "1.0.0", debug: bool = False):
        self.name = name
        self.version = version
        self.debug = debug

        # Initialize FastMCP (which creates the underlying Server)
        self.fastmcp = FastMCP(name, dependencies=[])

        # Access the underlying Server instance for the transport layer
        self.mcp_server = self.fastmcp._mcp_server

        # Initialize the SSE Transport
        self.sse_transport = SseServerTransport("/messages")

        self.router = APIRouter()
        self._setup_routes()

    def _setup_routes(self):
        """Register the SSE and Messages endpoints."""

        @self.router.get("/sse")
        async def handle_sse(request: Request):
            async with self.sse_transport.connect_sse(
                request.scope, request.receive, request._send
            ) as streams:
                read_stream, write_stream = streams

                # FastMCP sets up capabilities on the server automatically
                init_options = self.mcp_server.create_initialization_options()

                await self.mcp_server.run(
                    read_stream, write_stream, init_options, raise_exceptions=self.debug
                )

        @self.router.post("/messages")
        async def handle_messages(request: Request):
            return MCPAsgiResponse(self.sse_transport)

    # -------------------------------------------------------------------------
    # Delegated Decorators
    # -------------------------------------------------------------------------

    def tool(self, name: Optional[str] = None, **kwargs):
        return self.fastmcp.tool(name, **kwargs)

    def resource(self, uri: str, **kwargs):
        return self.fastmcp.resource(uri, **kwargs)

    # -------------------------------------------------------------------------
    # Auto-Registration Helpers
    # -------------------------------------------------------------------------

    def register_cube_resource(self, cube_cls: Any, uri: Optional[str] = None):
        """
        Automatically register a Datacube class as an MCP resource.

        The resource will be available at `sibi://<cube_name>` (concatenated).
        When read, it triggers `cube.aload()` and returns the data as JSON.
        """
        # Default URI: sibi://MyCubeName
        if not uri:
            uri = f"sibi://{cube_cls.__name__}"

        @self.resource(uri)
        async def handle_resource_read() -> str:
            # Instantiate and load
            # Note: We use the default config. If the cube needs specific kwargs,
            # this basic bridge might not be enough (would need a factory).
            cube = cube_cls()
            try:
                # Load data (async)
                df = await cube.aload()

                # Materialize (handling Dask)
                if hasattr(df, "compute"):
                    # Simple compute for now. Ideally we use dask_cluster.safe_compute
                    # but we keep imports minimal here to avoid circular dep risks in boilerplate.
                    # We can do a local import if needed or simple awaitable.
                    import asyncio

                    df = await asyncio.to_thread(df.compute)

                if df is None:
                    return "[]"

                # Return JSON string
                return df.to_json(orient="records", date_format="iso")

            except Exception as e:
                # Log the actual error
                print(f"MCP RESOURCE LOAD ERROR ({uri}): {str(e)}")
                # Return sanitized error to client
                return "Error loading resource: An internal error occurred."

        return handle_resource_read
