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

import json
import os
from typing import Any

from fastmcp import Client
from fastmcp.client import SSETransport, StdioTransport, StreamableHttpTransport
from fastmcp.client.client import CallToolResult

from .base_client import BaseClient
from ..component_registry import R
from ...constants import REME_SERVICE_INFO, REME_DEFAULT_HOST, REME_DEFAULT_PORT

_TRANSPORT_MAP = {
    "sse": SSETransport,
    "stdio": StdioTransport,
    "streamable-http": StreamableHttpTransport,
}


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

    Usage:
        # SSE (default)
        client = MCPClient(action="my_tool", host="localhost", port=8000, query="hello")
        async with client:
            result = await client()

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

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

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

    def __init__(
        self,
        action: str,
        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 _TRANSPORT_MAP:
            raise ValueError(f"Unknown transport: {transport!r}, expected one of {list(_TRANSPORT_MAP)}")

        if isinstance(transport, str) and transport != "stdio":
            if not (host and port):
                if service_info := os.environ.get(REME_SERVICE_INFO):
                    try:
                        data = json.loads(service_info)
                        host = data["host"]
                        port = data["port"]
                    except Exception:
                        self.logger.warning(f"Invalid service info: {service_info}")
                        host, port = REME_DEFAULT_HOST, REME_DEFAULT_PORT
                else:
                    host, port = REME_DEFAULT_HOST, REME_DEFAULT_PORT
            self.host = host
            self.port = port

        self.action = action
        self.transport = transport
        self.timeout = timeout

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

        cls = _TRANSPORT_MAP[self.transport]

        if self.transport == "stdio":
            command = self.kwargs.pop("command", "")
            args = self.kwargs.pop("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:
            self.client = Client(self._build_transport(), timeout=self.timeout)
            await self.client.__aenter__()

    async def __call__(self) -> str:
        if self.client is None:
            raise RuntimeError("Client not initialized. Call _start() first.")

        result: CallToolResult = await self.client.call_tool(self.action, self.kwargs)
        return self._extract_text(result)

    # 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)