"""MCP client for ReMe services."""

from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any

from .base_client import BaseClient
from ..component_registry import R

if TYPE_CHECKING:
    from fastmcp.client.client import CallToolResult

_VALID_TRANSPORTS = {"sse", "stdio", "streamable-http"}


@R.register("mcp")
class MCPClient(BaseClient):
    """MCP client that communicates with ReMe MCP service via fastmcp.Client.

    Usage:
        # SSE (default)
        client = MCPClient(host="localhost", port=8000)
        async with client:
            async for text in client(action="my_tool", query="hello"):
                print(text)

        # Streamable HTTP
        client = MCPClient(transport="streamable-http", host="localhost", port=8000)

        # Stdio
        client = MCPClient(transport="stdio", command="python", args=["server.py"])

        # Custom transport object
        from fastmcp.client import SSETransport
        client = MCPClient(transport=SSETransport(url="http://host:port/sse"))
    """

    def __init__(
        self,
        transport: str | Any = "sse",
        host: str | None = None,
        port: int | None = None,
        timeout: float = 30.0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if isinstance(transport, str) and transport not in _VALID_TRANSPORTS:
            raise ValueError(f"Unknown transport: {transport!r}, expected one of {sorted(_VALID_TRANSPORTS)}")

        if isinstance(transport, str) and transport != "stdio":
            self.host, self.port = self._resolve_service_address(host, port)

        self.transport = transport
        self.timeout = timeout

    def _build_transport(self):
        if not isinstance(self.transport, str):
            return self.transport

        from fastmcp.client import SSETransport, StdioTransport, StreamableHttpTransport

        transport_map = {
            "sse": SSETransport,
            "stdio": StdioTransport,
            "streamable-http": StreamableHttpTransport,
        }
        cls = transport_map[self.transport]

        if self.transport == "stdio":
            command = self.kwargs.get("command", "")
            args = self.kwargs.get("args", [])
            return cls(command=command, args=args)

        path = "/sse" if self.transport == "sse" else "/mcp"
        url = f"http://{self.host}:{self.port}{path}"
        return cls(url=url)

    # pylint: disable=unnecessary-dunder-call
    async def _start(self) -> None:
        if self.client is None:
            from fastmcp import Client

            self.client = Client(self._build_transport(), timeout=self.timeout)
            await self.client.__aenter__()

    # pylint: disable=invalid-overridden-method
    async def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]:
        if self.client is None:
            raise RuntimeError("Client not initialized. Call _start() first.")

        result: "CallToolResult" = await self.client.call_tool(action, payload)
        yield self._extract_text(result)

    async def list_actions(self) -> list[dict]:
        """Return raw MCP Tool dumps; each dict gets an `action` key (the tool name)."""
        if self.client is None:
            raise RuntimeError("Client not initialized. Call _start() first.")
        tools = await self.client.list_tools()
        return [tool.model_dump() for tool in tools]

    # pylint: disable=unnecessary-dunder-call
    async def _close(self) -> None:
        if self.client is not None:
            await self.client.__aexit__(None, None, None)
            self.client = None

    @staticmethod
    def _extract_text(result: "CallToolResult") -> str:
        for block in result.content:
            if hasattr(block, "text"):
                return block.text
        return str(result.content)