"""MCP client — connects to configured MCP servers and registers their tools."""

import json
from pathlib import Path
from typing import Any

import structlog

log = structlog.get_logger()


def load_mcp_config(project_root: Path) -> dict:
    """Load .fliiq/mcp.json. Returns empty dict if missing."""
    config_path = project_root / ".fliiq" / "mcp.json"
    if not config_path.exists():
        return {}
    try:
        return json.loads(config_path.read_text())
    except (json.JSONDecodeError, OSError) as e:
        log.warning("mcp_config_invalid", path=str(config_path), error=str(e))
        return {}


def save_mcp_config(project_root: Path, config: dict) -> None:
    """Write .fliiq/mcp.json."""
    config_path = project_root / ".fliiq" / "mcp.json"
    config_path.parent.mkdir(parents=True, exist_ok=True)
    config_path.write_text(json.dumps(config, indent=2) + "\n")


class MCPManager:
    """Manages connections to MCP servers and exposes their tools."""

    def __init__(self):
        self._sessions: dict[str, Any] = {}  # server_name -> (session, cleanup_cms)
        self._tools: dict[str, tuple[str, Any, str]] = {}  # tool_name -> (server_name, session, original_name)

    async def connect_all(self, config: dict, registry) -> list[dict]:
        """Connect to all configured MCP servers and register their tools.

        Returns list of {name, description} for prompt display.
        """
        servers = config.get("servers", {})
        if not servers:
            return []

        tool_infos = []

        for server_name, server_config in servers.items():
            try:
                session, cleanup = await self._connect_server(server_name, server_config)
                self._sessions[server_name] = (session, cleanup)

                # List tools from this server
                tools_response = await session.list_tools()
                for tool in tools_response.tools:
                    tool_name = f"mcp_{server_name}_{tool.name}"

                    # Capture per-iteration to avoid closure issues
                    _session = session
                    _original_name = tool.name

                    async def _handler(params: dict, s=_session, n=_original_name) -> str:
                        return await self._call_tool(s, n, params)

                    description = f"[MCP:{server_name}] {tool.description or tool.name}"
                    input_schema = tool.inputSchema if tool.inputSchema else {"type": "object", "properties": {}}

                    registry.register(
                        name=tool_name,
                        description=description,
                        input_schema=input_schema,
                        handler=_handler,
                    )
                    self._tools[tool_name] = (server_name, session, tool.name)
                    tool_infos.append({"name": tool_name, "description": description})

                log.info("mcp_server_connected", server=server_name, tools=len(tools_response.tools))

            except Exception as e:
                log.warning("mcp_server_failed", server=server_name, error=str(e))
                # Non-fatal: continue with other servers

        return tool_infos

    async def _connect_server(self, name: str, config: dict) -> tuple[Any, list]:
        """Connect to a single MCP server. Returns (session, list_of_context_managers)."""
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client
        from mcp.client.streamable_http import streamablehttp_client

        transport = config.get("transport", "stdio")
        cleanup_cms = []

        if transport == "streamable-http":
            url = config["url"]
            headers = config.get("headers")
            cm = streamablehttp_client(url=url, headers=headers)
            read_stream, write_stream, _ = await cm.__aenter__()
            cleanup_cms.append(cm)
        elif transport == "stdio":
            params = StdioServerParameters(
                command=config["command"],
                args=config.get("args", []),
                env=config.get("env"),
            )
            cm = stdio_client(params)
            read_stream, write_stream = await cm.__aenter__()
            cleanup_cms.append(cm)
        else:
            raise ValueError(f"Unknown MCP transport: {transport}")

        session_cm = ClientSession(read_stream, write_stream)
        session = await session_cm.__aenter__()
        cleanup_cms.insert(0, session_cm)  # Session closed first
        await session.initialize()

        return session, cleanup_cms

    async def _call_tool(self, session: Any, tool_name: str, params: dict) -> str:
        """Forward a tool call to an MCP server."""
        result = await session.call_tool(tool_name, arguments=params)

        if result.isError:
            parts = []
            for content in result.content:
                if hasattr(content, "text"):
                    parts.append(content.text)
                else:
                    parts.append(str(content))
            return f"MCP tool error: {' '.join(parts)}"

        parts = []
        for content in result.content:
            if hasattr(content, "text"):
                parts.append(content.text)
            else:
                parts.append(str(content))

        return "\n".join(parts) if parts else "Tool returned no content."

    async def disconnect_all(self) -> None:
        """Clean up all MCP connections."""
        for server_name, (session, cleanup_cms) in self._sessions.items():
            for cm in cleanup_cms:
                try:
                    await cm.__aexit__(None, None, None)
                except Exception as e:
                    log.debug("mcp_disconnect_error", server=server_name, error=str(e))
        self._sessions.clear()
        self._tools.clear()

    async def test_server(self, name: str, config: dict) -> dict:
        """Test a single MCP server connection. Returns {status, tools, error}."""
        try:
            session, cleanup = await self._connect_server(name, config)
            tools_response = await session.list_tools()
            tool_names = [t.name for t in tools_response.tools]

            # Cleanup
            for cm in cleanup:
                try:
                    await cm.__aexit__(None, None, None)
                except Exception:
                    pass

            return {"status": "ok", "tools": tool_names, "error": None}
        except Exception as e:
            return {"status": "error", "tools": [], "error": str(e)}
