"""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, Callable
from urllib.parse import urlparse

from ..config import McpServerConfig

# Type for status update callbacks: (server_name, status_dict) -> None
StatusCallback = Callable[[str, dict[str, Any]], None]

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

    @property
    def is_ready(self) -> bool:
        """True when all configured servers have resolved (connected or failed)."""
        if not self._configs:
            return True
        return len(self._server_status) >= len(self._configs)

    async def startup(self, status_callback: StatusCallback | None = None) -> 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) in parallel: %s", len(self._configs), ", ".join(self._configs))

        # Notify connecting status for all servers
        if status_callback:
            for name in self._configs:
                status_callback(name, {"status": "connecting", "tool_count": 0})

        # Connect all servers in parallel
        await asyncio.gather(
            *(self._connect_one(config, status_callback=status_callback) for config in self._configs.values()),
            return_exceptions=True,
        )

        # Rebuild tool map once after all connections complete (not per-server)
        self._rebuild_tool_map()

        # 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.info(
                "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, status_callback: StatusCallback | None = None) -> 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.warning("MCP server '%s': %s [%s]", config.name, error_msg, desc)
            self._server_status[config.name] = {
                "status": "error",
                "tool_count": 0,
                "error_message": error_msg,
            }

        # Notify caller of final status
        if status_callback:
            status_callback(config.name, self._server_status.get(config.name, {}))

    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.debug("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)
                errlog = open(os.devnull, "w")  # noqa: SIM115
                stack.callback(errlog.close)
                stdio_transport = await stack.enter_async_context(stdio_client(server_params, errlog=errlog))
                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.debug("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.debug("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

            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.warning(
                "MCP server '%s' failed to connect: %s [%s]",
                config.name,
                error_msg,
                desc,
                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])
        self._rebuild_tool_map()

    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

        # Resolve per-server trust level for context tagging
        srv_config = self._configs.get(server_name)
        trust = srv_config.trust_level if srv_config else "untrusted"

        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),
                "_context_trust": trust,
                "_context_origin": f"mcp:{server_name}",
            }

        return {"result": str(result), "_context_trust": trust, "_context_origin": f"mcp:{server_name}"}

    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
