"""In-memory transport for FastMCP servers — zero subprocess overhead."""

from __future__ import annotations

import json
import logging
from typing import Any

logger = logging.getLogger("mcp_forge.transport.inmemory")


class InMemoryTransport:
    """Connects directly to a FastMCP server instance in-process.

    This transport avoids spawning a subprocess by calling the FastMCP
    server's internal methods directly. Requires `fastmcp` to be installed.

    Usage:
        from my_server import mcp  # FastMCP instance
        transport = InMemoryTransport(mcp)
    """

    def __init__(self, server: Any) -> None:
        self._server = server
        self._client: Any = None

    async def start(self) -> None:
        """Initialize the in-memory client connection."""
        try:
            from fastmcp import Client
        except ImportError:
            raise ImportError(
                "FastMCP is required for in-memory transport. "
                "Install it with: pip install 'mcp-forge[fastmcp]'"
            ) from None

        self._client = Client(self._server)
        await self._client.__aenter__()

    async def send_request(
        self,
        method: str,
        params: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        """Route JSON-RPC method calls to the FastMCP client.

        This translates our generic JSON-RPC interface into FastMCP's
        typed client calls.
        """
        if self._client is None:
            raise ConnectionError("In-memory client not initialized")

        if method == "initialize":
            # FastMCP handles initialization internally
            return {
                "jsonrpc": "2.0",
                "id": 1,
                "result": {
                    "protocolVersion": "2025-06-18",
                    "capabilities": {
                        "tools": {"listChanged": False},
                        "resources": {"listChanged": False},
                        "prompts": {"listChanged": False},
                    },
                    "serverInfo": {"name": "fastmcp-inmemory", "version": "0.0.0"},
                },
            }

        elif method == "tools/list":
            tools = await self._client.list_tools()
            return {
                "jsonrpc": "2.0",
                "id": 1,
                "result": {
                    "tools": [
                        {
                            "name": t.name,
                            "description": getattr(t, "description", None),
                            "inputSchema": getattr(t, "inputSchema", {}),
                        }
                        for t in tools
                    ]
                },
            }

        elif method == "tools/call":
            params = params or {}
            name = params.get("name", "")
            arguments = params.get("arguments", {})
            try:
                result = await self._client.call_tool(name, arguments)
                # FastMCP returns different result types depending on version
                # Normalize to our expected format
                if hasattr(result, "content"):
                    # Newer FastMCP versions
                    content = []
                    for item in result.content:
                        content.append({
                            "type": getattr(item, "type", "text"),
                            "text": getattr(item, "text", str(item)),
                        })
                    return {
                        "jsonrpc": "2.0",
                        "id": 1,
                        "result": {
                            "content": content,
                            "isError": getattr(result, "isError", False),
                        },
                    }
                elif isinstance(result, list):
                    # Older FastMCP versions return list of content blocks
                    content = []
                    for item in result:
                        content.append({
                            "type": getattr(item, "type", "text"),
                            "text": getattr(item, "text", str(item)),
                        })
                    return {
                        "jsonrpc": "2.0",
                        "id": 1,
                        "result": {"content": content, "isError": False},
                    }
                else:
                    # Fallback: treat as text
                    return {
                        "jsonrpc": "2.0",
                        "id": 1,
                        "result": {
                            "content": [{"type": "text", "text": str(result)}],
                            "isError": False,
                        },
                    }
            except Exception as exc:
                return {
                    "jsonrpc": "2.0",
                    "id": 1,
                    "result": {
                        "content": [{"type": "text", "text": str(exc)}],
                        "isError": True,
                    },
                }

        elif method == "resources/list":
            resources = await self._client.list_resources()
            return {
                "jsonrpc": "2.0",
                "id": 1,
                "result": {
                    "resources": [
                        {
                            "uri": str(getattr(r, "uri", r)),
                            "name": getattr(r, "name", None),
                            "description": getattr(r, "description", None),
                            "mimeType": getattr(r, "mimeType", None),
                        }
                        for r in resources
                    ]
                },
            }

        elif method == "resources/read":
            params = params or {}
            uri = params.get("uri", "")
            content = await self._client.read_resource(uri)
            if isinstance(content, str):
                return {
                    "jsonrpc": "2.0",
                    "id": 1,
                    "result": {
                        "contents": [{"uri": uri, "text": content}]
                    },
                }
            return {
                "jsonrpc": "2.0",
                "id": 1,
                "result": {"contents": [{"uri": uri, "text": str(content)}]},
            }

        elif method == "prompts/list":
            prompts = await self._client.list_prompts()
            return {
                "jsonrpc": "2.0",
                "id": 1,
                "result": {
                    "prompts": [
                        {
                            "name": p.name,
                            "description": getattr(p, "description", None),
                            "arguments": getattr(p, "arguments", []),
                        }
                        for p in prompts
                    ]
                },
            }

        elif method == "prompts/get":
            params = params or {}
            name = params.get("name", "")
            arguments = params.get("arguments", {})
            result = await self._client.get_prompt(name, arguments)
            return {
                "jsonrpc": "2.0",
                "id": 1,
                "result": {
                    "messages": getattr(result, "messages", []),
                    "description": getattr(result, "description", None),
                },
            }

        else:
            raise RuntimeError(f"Unsupported method for in-memory transport: {method}")

    async def send_notification(
        self, method: str, params: dict[str, Any] | None = None
    ) -> None:
        """Notifications are no-ops for in-memory transport."""
        pass

    async def stop(self) -> None:
        """Close the in-memory client connection."""
        if self._client is not None:
            await self._client.__aexit__(None, None, None)
            self._client = None

    @property
    def is_alive(self) -> bool:
        return self._client is not None

    async def __aenter__(self) -> InMemoryTransport:
        await self.start()
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.stop()
