"""MCP client lifecycle and tool routing manager."""

from __future__ import annotations

import asyncio
import fnmatch
import ipaddress
import logging
import os
import shutil
import socket
from contextlib import AsyncExitStack
from typing import Any
from urllib.parse import urlparse

from ..config import McpServerConfig

logger = logging.getLogger(__name__)

_BLOCKED_NETWORKS = [
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("169.254.0.0/16"),
    ipaddress.ip_network("::1/128"),
    ipaddress.ip_network("fc00::/7"),
    ipaddress.ip_network("fe80::/10"),
]


def _validate_sse_url(url: str) -> None:
    """Block SSE URLs pointing to internal/metadata endpoints (with DNS resolution)."""
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise ValueError(f"Unsupported URL scheme: {parsed.scheme}")
    hostname = parsed.hostname or ""
    if hostname in ("localhost", "metadata.google.internal"):
        raise ValueError(f"Blocked internal hostname: {hostname}")
    try:
        addr = ipaddress.ip_address(hostname)
        for network in _BLOCKED_NETWORKS:
            if addr in network:
                raise ValueError(f"Blocked internal IP: {hostname}")
    except ValueError as e:
        if "Blocked" in str(e):
            raise
        try:
            infos = socket.getaddrinfo(hostname, None, proto=socket.IPPROTO_TCP)
            for _family, _type, _proto, _canon, sockaddr in infos:
                resolved_addr = ipaddress.ip_address(sockaddr[0])
                for network in _BLOCKED_NETWORKS:
                    if resolved_addr in network:
                        raise ValueError(f"Hostname '{hostname}' resolves to blocked IP: {sockaddr[0]}")
        except socket.gaierror:
            raise ValueError(f"Cannot resolve hostname: {hostname}")


def _validate_command(command: str) -> None:
    """Validate MCP command exists on PATH."""
    resolved = shutil.which(command)
    if resolved is None:
        raise ValueError(f"MCP command not found on PATH: {command}")


class McpManager:
    def __init__(self, server_configs: list[McpServerConfig], tool_warning_threshold: int = 40) -> None:
        self._configs: dict[str, McpServerConfig] = {cfg.name: cfg for cfg in server_configs}
        self._exit_stacks: dict[str, AsyncExitStack] = {}
        self._sessions: dict[str, Any] = {}
        self._server_tools: dict[str, list[dict[str, Any]]] = {}
        self._tool_to_server: dict[str, str] = {}
        self._server_status: dict[str, dict[str, Any]] = {}
        self._disabled: set[str] = set()
        self._tool_warning_threshold = tool_warning_threshold

    async def startup(self) -> None:
        if not self._configs:
            return

        try:
            from mcp import ClientSession, StdioServerParameters  # noqa: F401
            from mcp.client.stdio import stdio_client  # noqa: F401
        except ImportError:
            logger.warning("MCP SDK not installed, skipping MCP server connections")
            return

        logger.info("Starting %d MCP server(s): %s", len(self._configs), ", ".join(self._configs))
        for name, config in self._configs.items():
            await self._connect_one(config)

        # Log startup summary
        connected = [n for n, s in self._server_status.items() if s.get("status") == "connected"]
        failed = [n for n, s in self._server_status.items() if s.get("status") == "error"]
        if failed:
            logger.warning(
                "MCP startup complete: %d connected, %d failed (%s)",
                len(connected),
                len(failed),
                ", ".join(failed),
            )
        else:
            logger.info("MCP startup complete: %d/%d connected", len(connected), len(self._configs))

        # Warn if total tool count exceeds threshold
        total_tools = len(self._tool_to_server)
        if self._tool_warning_threshold > 0 and total_tools > self._tool_warning_threshold:
            per_server = []
            any_filters_active = False
            for sname in connected:
                cfg = self._configs.get(sname)
                has_filter = bool(cfg and (cfg.tools_include or cfg.tools_exclude)) if cfg else False
                if has_filter:
                    any_filters_active = True
                count = sum(1 for t in self._server_tools.get(sname, []) if self._is_tool_allowed(t["name"], sname))
                per_server.append(f"{sname}: {count}")
            hint = (
                "Current filters may not be restrictive enough."
                if any_filters_active
                else "Consider adding tools_include/tools_exclude filters to your MCP server config."
            )
            logger.warning(
                "MCP tool count (%d) exceeds threshold (%d). %s Per-server: %s",
                total_tools,
                self._tool_warning_threshold,
                hint,
                ", ".join(per_server),
            )

    def _describe_config(self, config: McpServerConfig) -> str:
        """Build a human-readable description of a server config for diagnostics."""
        parts = [f"name={config.name}", f"transport={config.transport}"]
        if config.command:
            cmd = f"{config.command} {' '.join(config.args)}" if config.args else config.command
            parts.append(f"command={cmd}")
        if config.url:
            parts.append(f"url={config.url}")
        if config.env:
            parts.append(f"env_keys={list(config.env.keys())}")
        parts.append(f"timeout={config.timeout}s")
        return ", ".join(parts)

    async def _connect_one(self, config: McpServerConfig) -> None:
        desc = self._describe_config(config)
        logger.info("Connecting MCP server '%s' [%s] (timeout=%gs)", config.name, desc, config.timeout)

        try:
            await asyncio.wait_for(self._do_connect(config), timeout=config.timeout)
        except asyncio.TimeoutError:
            error_msg = f"Connection timed out after {config.timeout}s"
            logger.error("MCP server '%s': %s [%s]", config.name, error_msg, desc)
            self._server_status[config.name] = {
                "status": "error",
                "tool_count": 0,
                "error_message": error_msg,
            }

    async def _do_connect(self, config: McpServerConfig) -> None:
        """Inner connection logic, called under a timeout by _connect_one."""
        desc = self._describe_config(config)

        try:
            from mcp import ClientSession, StdioServerParameters
            from mcp.client.stdio import stdio_client
        except ImportError:
            msg = "MCP SDK not installed (pip install mcp)"
            logger.error("MCP server '%s': %s", config.name, msg)
            self._server_status[config.name] = {
                "status": "error",
                "tool_count": 0,
                "error_message": msg,
            }
            return

        stack = AsyncExitStack()
        try:
            if config.transport == "stdio" and config.command:
                resolved = shutil.which(config.command)
                if resolved is None:
                    raise FileNotFoundError(
                        f"Command '{config.command}' not found on PATH. Searched: {os.environ.get('PATH', '(empty)')}"
                    )
                logger.debug(
                    "MCP '%s': command '%s' resolved to '%s'",
                    config.name,
                    config.command,
                    resolved,
                )
                server_params = StdioServerParameters(
                    command=config.command,
                    args=config.args,
                    env={**os.environ, **config.env} if config.env else None,
                )
                logger.debug("MCP '%s': launching stdio transport", config.name)
                stdio_transport = await stack.enter_async_context(stdio_client(server_params))
                read_stream, write_stream = stdio_transport
                logger.debug("MCP '%s': initializing session", config.name)
                session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
                await session.initialize()
                logger.debug("MCP '%s': session initialized", config.name)

            elif config.transport == "sse" and config.url:
                try:
                    from mcp.client.sse import sse_client
                except ImportError:
                    msg = "SSE client not available (pip install mcp[sse])"
                    logger.error("MCP server '%s': %s", config.name, msg)
                    self._server_status[config.name] = {
                        "status": "error",
                        "tool_count": 0,
                        "error_message": msg,
                    }
                    return

                _validate_sse_url(config.url)
                logger.debug("MCP '%s': connecting SSE to %s", config.name, config.url)
                sse_transport = await stack.enter_async_context(sse_client(config.url))
                read_stream, write_stream = sse_transport
                logger.debug("MCP '%s': initializing session", config.name)
                session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
                await session.initialize()
                logger.debug("MCP '%s': session initialized", config.name)

            else:
                msg = (
                    f"Invalid transport config: transport='{config.transport}', "
                    f"command={'set' if config.command else 'missing'}, "
                    f"url={'set' if config.url else 'missing'}"
                )
                logger.error("MCP server '%s': %s", config.name, msg)
                self._server_status[config.name] = {
                    "status": "error",
                    "tool_count": 0,
                    "error_message": msg,
                }
                return

            self._exit_stacks[config.name] = stack
            self._sessions[config.name] = session

            logger.debug("MCP '%s': listing tools", config.name)
            tools_result = await session.list_tools()
            server_tools: list[dict[str, Any]] = []
            for tool in tools_result.tools:
                tool_entry = {
                    "name": tool.name,
                    "server_name": config.name,
                    "description": tool.description or "",
                    "input_schema": tool.inputSchema if hasattr(tool, "inputSchema") else {},
                }
                server_tools.append(tool_entry)

            self._server_tools[config.name] = server_tools
            self._rebuild_tool_map()

            total = len(server_tools)
            allowed_tools = [t for t in server_tools if self._is_tool_allowed(t["name"], config.name)]
            allowed_count = len(allowed_tools)
            tool_names = [t["name"] for t in allowed_tools]
            self._server_status[config.name] = {
                "status": "connected",
                "tool_count": allowed_count,
                "total_tool_count": total,
            }
            if allowed_count < total:
                logger.info(
                    "MCP server '%s' connected: %d/%d tools (filtered) (%s)",
                    config.name,
                    allowed_count,
                    total,
                    ", ".join(tool_names[:10]) + ("..." if len(tool_names) > 10 else ""),
                )
            else:
                logger.info(
                    "MCP server '%s' connected: %d tools (%s)",
                    config.name,
                    total,
                    ", ".join(tool_names[:10]) + ("..." if len(tool_names) > 10 else ""),
                )

        except BaseException as e:
            # Close the exit stack to clean up any partially-entered async
            # contexts (stdio subprocess, task groups, etc.) so they don't
            # leak and crash with "unhandled exception in a TaskGroup".
            try:
                await stack.aclose()
            except BaseException:
                logger.debug(
                    "Error closing stack for '%s' during cleanup",
                    config.name,
                    exc_info=True,
                )

            # Re-raise CancelledError so asyncio.wait_for's timeout mechanism
            # functions correctly across Python versions.
            if isinstance(e, asyncio.CancelledError):
                raise

            # Import McpError here (MCP SDK may not be installed).
            try:
                from mcp import McpError
            except ImportError:
                McpError = None  # type: ignore[assignment,misc]  # noqa: N806

            is_mcp_error = McpError is not None and isinstance(e, McpError)
            error_type = type(e).__name__
            error_msg = f"{error_type}: {e}"
            logger.error(
                "MCP server '%s' failed to connect: %s [%s]",
                config.name,
                error_msg,
                desc,
                # Suppress full traceback for known MCP protocol errors — the
                # message is already diagnostic.  Keep it for unexpected errors.
                exc_info=not is_mcp_error,
            )
            self._server_status[config.name] = {
                "status": "error",
                "tool_count": 0,
                "error_message": error_msg,
            }

    def _is_tool_allowed(self, tool_name: str, server_name: str) -> bool:
        """Check if a tool passes the server's include/exclude filter."""
        config = self._configs.get(server_name)
        if not config:
            return True
        if config.tools_include:
            return any(fnmatch.fnmatch(tool_name, pattern) for pattern in config.tools_include)
        if config.tools_exclude:
            return not any(fnmatch.fnmatch(tool_name, pattern) for pattern in config.tools_exclude)
        return True

    def _rebuild_tool_map(self) -> None:
        """Rebuild _tool_to_server from _server_tools, warning on collisions."""
        self._tool_to_server.clear()
        for server_name, tools in self._server_tools.items():
            for tool in tools:
                if not self._is_tool_allowed(tool["name"], server_name):
                    continue
                if tool["name"] in self._tool_to_server:
                    logger.warning(
                        "Tool name collision: '%s' from server '%s' shadows existing tool from '%s'",
                        tool["name"],
                        server_name,
                        self._tool_to_server[tool["name"]],
                    )
                self._tool_to_server[tool["name"]] = server_name

    async def connect_server(self, name: str) -> None:
        """Connect (or reconnect) a single server by name."""
        if name not in self._configs:
            raise ValueError(f"Unknown MCP server: {name}")

        # Disconnect first if already connected
        if name in self._sessions:
            await self.disconnect_server(name)

        self._disabled.discard(name)
        await self._connect_one(self._configs[name])

    async def disconnect_server(self, name: str) -> None:
        """Disconnect a single server by name."""
        if name not in self._configs:
            raise ValueError(f"Unknown MCP server: {name}")

        self._disabled.add(name)
        self._sessions.pop(name, None)
        self._server_tools.pop(name, None)
        self._rebuild_tool_map()

        stack = self._exit_stacks.pop(name, None)
        if stack:
            try:
                await stack.aclose()
            except BaseException:
                logger.warning(f"Error closing exit stack for '{name}'", exc_info=True)

        self._server_status[name] = {
            "status": "disconnected",
            "tool_count": 0,
        }
        logger.info(f"MCP server '{name}' disconnected")

    async def reconnect_server(self, name: str) -> None:
        """Disconnect then reconnect a server."""
        await self.connect_server(name)

    def get_openai_tools(self) -> list[dict[str, Any]] | None:
        all_tools = self.get_all_tools()
        if not all_tools:
            return None
        return [
            {
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": tool["input_schema"],
                },
            }
            for tool in all_tools
        ]

    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any:
        server_name = self._tool_to_server.get(tool_name)
        if not server_name or server_name not in self._sessions:
            # Check if this tool exists but was filtered out
            for sname, tools in self._server_tools.items():
                if any(t["name"] == tool_name for t in tools):
                    raise ValueError(f"MCP tool '{tool_name}' is filtered out by server '{sname}' tool configuration")
            available = list(self._sessions.keys())
            raise ValueError(
                f"MCP tool '{tool_name}' not available. "
                f"Connected servers: {', '.join(available) if available else '(none)'}"
            )

        # SECURITY-REVIEW: MCP tool arguments are JSON-serialized over stdio/SSE and never
        # passed to a shell. Input validation for shell injection belongs in tools/safety.py
        # (the bash tool boundary), not here. Removed _validate_tool_args per #291.
        session = self._sessions[server_name]
        try:
            result = await session.call_tool(tool_name, arguments)
        except Exception as e:
            logger.error("MCP tool '%s' on server '%s' failed: %s", tool_name, server_name, e, exc_info=True)
            raise ValueError(f"MCP server '{server_name}' returned an error for tool '{tool_name}'") from e

        if hasattr(result, "content"):
            contents = []
            for item in result.content:
                if hasattr(item, "text"):
                    contents.append(item.text)
                elif hasattr(item, "data"):
                    contents.append(str(item.data))
                else:
                    contents.append(str(item))
            return {"content": "\n".join(contents)}

        return {"result": str(result)}

    def get_tool_server_name(self, tool_name: str) -> str:
        return self._tool_to_server.get(tool_name, "unknown")

    def get_all_tools(self) -> list[dict[str, Any]]:
        """Flatten _server_tools into a single list, respecting per-server filters."""
        tools: list[dict[str, Any]] = []
        for server_name, server_tools in self._server_tools.items():
            for tool in server_tools:
                if self._is_tool_allowed(tool["name"], server_name):
                    tools.append(tool)
        return tools

    def get_server_statuses(self) -> dict[str, dict[str, Any]]:
        result = {}
        for name, config in self._configs.items():
            status = self._server_status.get(name, {"status": "disconnected", "tool_count": 0})
            result[name] = {
                "name": name,
                "transport": config.transport,
                **status,
            }
        return result

    async def shutdown(self) -> None:
        saved_interrupt: BaseException | None = None
        for name, stack in list(self._exit_stacks.items()):
            try:
                await asyncio.wait_for(stack.aclose(), timeout=5.0)
            except (KeyboardInterrupt, SystemExit) as e:
                logger.debug("Error closing exit stack for '%s' during shutdown", name, exc_info=True)
                if saved_interrupt is None:
                    saved_interrupt = e
            except BaseException:
                logger.debug("Error closing exit stack for '%s' during shutdown", name, exc_info=True)
        self._exit_stacks.clear()
        self._sessions.clear()
        self._server_tools.clear()
        self._tool_to_server.clear()
        if saved_interrupt is not None:
            raise saved_interrupt
