"""MCP client connector - handles connecting to MCP servers via stdio or HTTP."""

from __future__ import annotations

import asyncio
import logging
import os
import shutil
import subprocess
from contextlib import AsyncExitStack
from typing import Any

from mcp_bridge.registry.models import ServerEntry, ToolInfo, TransportType

logger = logging.getLogger(__name__)


class MCPConnection:
    """
    A connection to a single MCP server instance.

    Supports both STDIO and HTTP transports.
    Handles tool discovery, tool calling, and lifecycle management.
    """

    def __init__(self, server_entry: ServerEntry, instance_config: dict[str, Any]):
        self.server_entry = server_entry
        self.instance_config = instance_config
        self._process: subprocess.Popen | None = None
        self._exit_stack: AsyncExitStack | None = None
        self._session = None  # MCP ClientSession
        self._read_stream = None
        self._write_stream = None
        self._tools: list[ToolInfo] = []
        self._connected = False
        self._port: int | None = None
        self._http_session_id: str | None = None
        self._jsonrpc_counter: int = 0

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def pid(self) -> int | None:
        return self._process.pid if self._process else None

    @property
    def port(self) -> int | None:
        return self._port

    @property
    def tools(self) -> list[ToolInfo]:
        return self._tools

    async def connect(self) -> list[ToolInfo]:
        """Connect to the MCP server and discover tools."""
        if self.server_entry.transport == TransportType.STDIO:
            return await self._connect_stdio()
        elif self.server_entry.transport in (TransportType.HTTP, TransportType.SSE):
            return await self._connect_http()
        else:
            raise ValueError(f"Unsupported transport: {self.server_entry.transport}")

    async def _connect_stdio(self) -> list[ToolInfo]:
        """Connect via STDIO transport using the MCP SDK."""
        try:
            from mcp import ClientSession
            from mcp.client.stdio import StdioServerParameters, stdio_client

            # Build command
            command = self.server_entry.command
            args = list(self.server_entry.args)

            # Build environment -- start with OS env, layer server env, then
            # pass instance config as MCP_CFG_<KEY> env vars so servers can
            # optionally read them.  We do NOT blindly append config as CLI
            # flags because most MCP servers reject unknown arguments.
            env = dict(os.environ)
            env.update(self.server_entry.env)
            for key, value in self.instance_config.items():
                env[f"MCP_CFG_{key.upper()}"] = str(value)

            # Resolve command path
            cmd_parts = command.split()
            actual_command = cmd_parts[0]
            if len(cmd_parts) > 1:
                args = cmd_parts[1:] + args

            # For npx/node commands, find the full path
            cmd_path = shutil.which(actual_command)
            if cmd_path:
                actual_command = cmd_path

            server_params = StdioServerParameters(
                command=actual_command,
                args=args,
                env=env,
                cwd=self.server_entry.cwd,
            )

            self._exit_stack = AsyncExitStack()
            stdio_transport = await self._exit_stack.enter_async_context(stdio_client(server_params))
            self._read_stream, self._write_stream = stdio_transport

            self._session = await self._exit_stack.enter_async_context(
                ClientSession(self._read_stream, self._write_stream)
            )

            await self._session.initialize()

            # Discover tools
            tools_result = await self._session.list_tools()
            self._tools = []
            for tool in tools_result.tools:
                tool_info = ToolInfo(
                    name=tool.name,
                    description=tool.description or "",
                    input_schema=tool.inputSchema if hasattr(tool, "inputSchema") else {},
                    server_id=self.server_entry.id,
                )
                self._tools.append(tool_info)

            self._connected = True
            logger.info(
                "Connected to '%s' via stdio, discovered %d tools",
                self.server_entry.name,
                len(self._tools),
            )
            return self._tools

        except Exception as e:
            logger.error("Failed to connect to '%s' via stdio: %s", self.server_entry.name, e)
            await self.disconnect()
            raise

    def _next_jsonrpc_id(self) -> int:
        """Return a monotonically increasing JSON-RPC request id."""
        self._jsonrpc_counter += 1
        return self._jsonrpc_counter

    def _http_headers(self) -> dict[str, str]:
        """Build HTTP headers, including the MCP session ID when available."""
        headers: dict[str, str] = {}
        if self._http_session_id:
            headers["mcp-session-id"] = self._http_session_id
        return headers

    async def _connect_http(self) -> list[ToolInfo]:
        """Connect via HTTP transport - start server process and use JSON-RPC."""
        import aiohttp

        try:
            # Determine port
            self._port = self.server_entry.port or self._find_free_port()
            endpoint = self.server_entry.endpoint or "/mcp"

            # Build command
            cmd_parts = self.server_entry.command.split()
            full_args = cmd_parts + list(self.server_entry.args)

            # Build environment -- start with OS env, layer server env, then
            # pass instance config as MCP_CFG_<KEY> env vars (matching STDIO
            # behaviour) and expose the allocated port so the child process
            # knows which port to bind.
            env = dict(os.environ)
            env.update(self.server_entry.env)
            for key, value in self.instance_config.items():
                env[f"MCP_CFG_{key.upper()}"] = str(value)
            env["PORT"] = str(self._port)
            env["MCP_CFG_PORT"] = str(self._port)

            # Start the server process
            self._process = subprocess.Popen(
                full_args,
                env=env,
                cwd=self.server_entry.cwd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )

            # Wait for the server to be ready
            base_url = f"http://127.0.0.1:{self._port}"
            url = f"{base_url}{endpoint}"

            await self._wait_for_server(url)

            # Initialize MCP session
            async with aiohttp.ClientSession() as http_session:
                # Initialize
                init_response = await http_session.post(
                    url,
                    json={
                        "jsonrpc": "2.0",
                        "id": self._next_jsonrpc_id(),
                        "method": "initialize",
                        "params": {
                            "protocolVersion": "2024-11-05",
                            "capabilities": {},
                            "clientInfo": {"name": "mcp-bridge", "version": "1.0.0-beta-rc"},
                        },
                    },
                )
                await init_response.json()

                # Persist the session ID so all subsequent requests include it
                self._http_session_id = (
                    init_response.headers.get("mcp-session-id") or None
                )

                # List tools (include session header)
                tools_response = await http_session.post(
                    url,
                    json={
                        "jsonrpc": "2.0",
                        "id": self._next_jsonrpc_id(),
                        "method": "tools/list",
                        "params": {},
                    },
                    headers=self._http_headers(),
                )
                tools_data = await tools_response.json()

            # Parse tools
            self._tools = []
            if "result" in tools_data and "tools" in tools_data["result"]:
                for tool in tools_data["result"]["tools"]:
                    tool_info = ToolInfo(
                        name=tool["name"],
                        description=tool.get("description", ""),
                        input_schema=tool.get("inputSchema", {}),
                        server_id=self.server_entry.id,
                    )
                    self._tools.append(tool_info)

            self._connected = True
            logger.info(
                "Connected to '%s' via HTTP on port %d, discovered %d tools",
                self.server_entry.name,
                self._port,
                len(self._tools),
            )
            return self._tools

        except Exception as e:
            logger.error("Failed to connect to '%s' via HTTP: %s", self.server_entry.name, e)
            await self.disconnect()
            raise

    async def _wait_for_server(self, url: str, timeout: int = 30) -> None:
        """Wait for an HTTP server to become available."""
        import aiohttp

        start = asyncio.get_event_loop().time()
        while (asyncio.get_event_loop().time() - start) < timeout:
            try:
                async with aiohttp.ClientSession() as session, session.get(url, timeout=aiohttp.ClientTimeout(total=2)):
                    return
            except (TimeoutError, aiohttp.ClientError, ConnectionRefusedError):
                await asyncio.sleep(0.5)

        raise TimeoutError(f"Server at {url} did not become available within {timeout}s")

    def _find_free_port(self) -> int:
        """Find a free port on localhost."""
        import socket

        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("127.0.0.1", 0))
            return s.getsockname()[1]

    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any:
        """Call a tool on the connected server."""
        if not self._connected:
            raise RuntimeError("Not connected to server")

        if self.server_entry.transport == TransportType.STDIO:
            return await self._call_tool_stdio(tool_name, arguments)
        else:
            return await self._call_tool_http(tool_name, arguments)

    async def _call_tool_stdio(self, tool_name: str, arguments: dict[str, Any]) -> Any:
        """Call a tool via STDIO session."""
        assert self._session is not None
        result = await self._session.call_tool(tool_name, arguments)

        # Extract content from the result
        if hasattr(result, "content"):
            contents = []
            for item in result.content:
                if hasattr(item, "text"):
                    contents.append(item.text)
                elif hasattr(item, "data"):
                    contents.append(item.data)
                else:
                    contents.append(str(item))
            return contents[0] if len(contents) == 1 else contents

        return str(result)

    async def _call_tool_http(self, tool_name: str, arguments: dict[str, Any]) -> Any:
        """Call a tool via HTTP JSON-RPC."""
        import aiohttp

        endpoint = self.server_entry.endpoint or "/mcp"
        url = f"http://127.0.0.1:{self._port}{endpoint}"

        async with aiohttp.ClientSession() as session:
            response = await session.post(
                url,
                json={
                    "jsonrpc": "2.0",
                    "id": self._next_jsonrpc_id(),
                    "method": "tools/call",
                    "params": {"name": tool_name, "arguments": arguments},
                },
                headers=self._http_headers(),
            )
            data = await response.json()

        if "error" in data:
            raise RuntimeError(f"Tool call failed: {data['error']}")

        result = data.get("result", {})
        if "content" in result:
            contents = []
            for item in result["content"]:
                contents.append(item.get("text", str(item)))
            return contents[0] if len(contents) == 1 else contents

        return result

    async def disconnect(self) -> None:
        """Disconnect from the server and clean up resources."""
        self._connected = False
        self._tools = []
        self._http_session_id = None

        # Close STDIO session
        if self._exit_stack:
            try:
                await self._exit_stack.aclose()
            except Exception as e:
                logger.debug("Error closing exit stack: %s", e)
            self._exit_stack = None
            self._session = None

        # Kill HTTP process
        if self._process:
            try:
                self._process.terminate()
                try:
                    self._process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    self._process.kill()
                    self._process.wait(timeout=2)
            except Exception as e:
                logger.debug("Error terminating process: %s", e)
            self._process = None

        logger.info("Disconnected from '%s'", self.server_entry.name)

    async def health_check(self) -> bool:
        """Check if the connection is still alive."""
        if not self._connected:
            return False

        try:
            if self.server_entry.transport == TransportType.STDIO:
                # Try listing tools as a health check
                assert self._session is not None
                await self._session.list_tools()
                return True
            else:
                # Try an HTTP ping with session header
                import aiohttp

                endpoint = self.server_entry.endpoint or "/mcp"
                url = f"http://127.0.0.1:{self._port}{endpoint}"
                async with aiohttp.ClientSession(headers=self._http_headers()) as session, session.get(url, timeout=aiohttp.ClientTimeout(total=5)):
                    return True
        except Exception:
            return False
