"""MCP server connection manager for PentestAgent.

Uses standard MCP configuration format:
{
    "mcpServers": {
        "server-name": {
            "command": "npx",
            "args": ["-y", "package-name"],
            "env": {"VAR": "value"}
        }
    }
}
"""

import asyncio
import atexit
import json
import os
import signal
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional

from .tools import create_mcp_tool
from ..tools.registry import register_tool_instance, unregister_tool
from .transport import MCPTransport, StdioTransport


@dataclass
class MCPServerConfig:
    """Configuration for an MCP server."""

    name: str
    command: str
    args: List[str] = field(default_factory=list)
    env: Dict[str, str] = field(default_factory=dict)
    enabled: bool = True
    description: str = ""
    # Whether to auto-start this server when `connect_all()` is called.
    start_on_launch: bool = False


@dataclass
class MCPServer:
    """Represents a connected MCP server."""

    name: str
    config: MCPServerConfig
    transport: MCPTransport
    tools: List[dict] = field(default_factory=list)
    connected: bool = False
    # Lock for serializing all communication with this server
    # Prevents message ID collisions and transport interleaving
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)

    async def disconnect(self):
        """Disconnect from the server."""
        if self.connected:
            await self.transport.disconnect()
            self.connected = False


class MCPManager:
    """Manages MCP server connections and exposes tools to agents."""

    DEFAULT_CONFIG_PATHS = [
        Path.cwd() / "mcp_servers.json",
        Path.cwd() / "mcp.json",
        Path(__file__).parent / "mcp_servers.json",
        Path.home() / ".pentestagent" / "mcp_servers.json",
    ]

    def __init__(self, config_path: Optional[Path] = None):
        self.config_path = config_path or self._find_config()
        self.servers: Dict[str, MCPServer] = {}
        # Track adapters we auto-started so we can stop them later
        self._started_adapters: Dict[str, object] = {}
        self._message_id = 0
        # Control socket server attributes
        self._control_server: Optional[asyncio.AbstractServer] = None
        self._control_task: Optional[asyncio.Task] = None
        self._control_path: Optional[Path] = None
        # Auto-connect to configured MCP servers if any are enabled in config.
        try:
            cfg = self._load_config()
            has_enabled = any(getattr(s, "enabled", True) for s in cfg.values())
            if has_enabled:
                try:
                    loop = asyncio.get_running_loop()
                except RuntimeError:
                    # No running loop: run connect synchronously so tools register
                    try:
                        asyncio.run(self.connect_all())
                    except Exception:
                        pass
                else:
                    # Running loop present: schedule background connect task
                    try:
                        loop.create_task(self.connect_all())
                    except Exception:
                        pass
        except Exception:
            pass
        # Ensure we attempt to clean up vendored servers on process exit
        try:
            atexit.register(self._atexit_cleanup)
        except Exception as e:
            logging.getLogger(__name__).exception("Failed to register atexit cleanup: %s", e)
            try:
                from ..interface.notifier import notify

                notify("warning", f"Failed to register MCP atexit cleanup: {e}")
            except Exception:
                logging.getLogger(__name__).exception("Failed to notify operator about atexit.register failure")

    def _find_config(self) -> Path:
        for path in self.DEFAULT_CONFIG_PATHS:
            if path.exists():
                return path
        return self.DEFAULT_CONFIG_PATHS[0]

    def _get_next_id(self) -> int:
        self._message_id += 1
        return self._message_id

    def _load_config(self) -> Dict[str, MCPServerConfig]:
        if not self.config_path.exists():
            return {}
        try:
            raw = json.loads(self.config_path.read_text(encoding="utf-8"))
            servers = {}
            mcp_servers = raw.get("mcpServers", {})
            for name, config in mcp_servers.items():
                if not config.get("command"):
                    continue
                servers[name] = MCPServerConfig(
                    name=name,
                    command=config["command"],
                    args=config.get("args", []),
                    env=config.get("env", {}),
                    enabled=config.get("enabled", True),
                    start_on_launch=config.get("start_on_launch", False),
                    description=config.get("description", ""),
                )
            # Environment-based auto-start overrides (previously supported via
            # LAUNCH_* variables) have been removed. MCP adapters are no
            # longer auto-started by the manager; operators should install and
            # run adapters manually and add them to `mcp_servers.json` if they
            # want `MCPManager` to connect to them.

            return servers
        except json.JSONDecodeError as e:
            print(f"[MCP] Error loading config: {e}")
            return {}

    def _save_config(self, servers: Dict[str, MCPServerConfig]):
        config = {"mcpServers": {}}
        for name, server in servers.items():
            server_config = {"command": server.command, "args": server.args}
            if server.env:
                server_config["env"] = server.env
            if server.description:
                server_config["description"] = server.description
            if not server.enabled:
                server_config["enabled"] = False
            config["mcpServers"][name] = server_config
        self.config_path.parent.mkdir(parents=True, exist_ok=True)
        self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")

    def _atexit_cleanup(self):
        """Synchronous atexit cleanup that attempts to stop adapters and disconnect servers."""
        try:
            # Try to run async shutdown; if an event loop is already running this may fail,
            # but it's best-effort to avoid orphaned vendored servers.
            asyncio.run(self._stop_started_adapters_and_disconnect())
        except Exception:
            # Last-ditch: attempt to stop adapters synchronously.
            # If the adapter exposes a blocking `stop()` call, call it. Otherwise, try
            # to kill the underlying process by PID to avoid asyncio subprocess
            # destructors running after the loop is closed.
            for adapter in list(self._started_adapters.values()):
                try:
                    # Prefer adapter-provided synchronous stop hook
                    stop_sync = getattr(adapter, "stop_sync", None)
                    if stop_sync:
                        try:
                            stop_sync()
                            continue
                        except Exception:
                            pass

                    # Fallback: try blocking stop() if present
                    stop = getattr(adapter, "stop", None)
                    if stop and not asyncio.iscoroutinefunction(stop):
                        try:
                            stop()
                            continue
                        except Exception as e:
                            logging.getLogger(__name__).exception(
                                "Error running adapter.stop(): %s", e
                            )

                    # Final fallback: kill underlying PID if available
                    pid = None
                    proc = getattr(adapter, "_process", None)
                    if proc is not None:
                        pid = getattr(proc, "pid", None)
                    if pid:
                        try:
                            os.kill(pid, signal.SIGTERM)
                        except Exception as e:
                            logging.getLogger(__name__).exception("Failed to SIGTERM pid %s: %s", pid, e)
                            try:
                                os.kill(pid, signal.SIGKILL)
                            except Exception as e2:
                                logging.getLogger(__name__).exception("Failed to SIGKILL pid %s: %s", pid, e2)
                except Exception as e:
                    logging.getLogger(__name__).exception("Error while attempting synchronous adapter stop: %s", e)

    async def _stop_started_adapters_and_disconnect(self) -> None:
        # Stop any adapters we started
        for _name, adapter in list(self._started_adapters.items()):
            try:
                stop = getattr(adapter, "stop", None)
                if stop:
                    if asyncio.iscoroutinefunction(stop):
                        await stop()
                    else:
                        # run blocking stop in executor
                        loop = asyncio.get_running_loop()
                        await loop.run_in_executor(None, stop)
            except Exception as e:
                logging.getLogger(__name__).exception("Error stopping adapter in async shutdown: %s", e)
        self._started_adapters.clear()

        # Disconnect any active MCP server connections
        try:
            await self.disconnect_all()
        except Exception as e:
            logging.getLogger(__name__).exception("Error during disconnect_all in shutdown: %s", e)

    def add_server(
        self,
        name: str,
        command: str,
        args: List[str] = None,
        env: Dict[str, str] = None,
        description: str = "",
    ):
        servers = self._load_config()
        servers[name] = MCPServerConfig(
            name=name,
            command=command,
            args=args or [],
            env=env or {},
            description=description,
        )
        self._save_config(servers)
        print(f"[MCP] Added server: {name}")

    def remove_server(self, name: str) -> bool:
        servers = self._load_config()
        if name in servers:
            del servers[name]
            self._save_config(servers)
            return True
        return False

    def set_enabled(self, name: str, enabled: bool) -> bool:
        """Enable or disable a configured MCP server in the config file.

        Returns True if the server existed and was updated, False otherwise.
        """
        servers = self._load_config()
        if name not in servers:
            return False
        servers[name].enabled = bool(enabled)
        self._save_config(servers)
        return True

    def list_configured_servers(self) -> List[dict]:
        servers = self._load_config()
        return [
            {
                "name": n,
                "command": s.command,
                "args": s.args,
                "env": s.env,
                "enabled": s.enabled,
                "description": s.description,
                "connected": n in self.servers and self.servers[n].connected,
            }
            for n, s in servers.items()
        ]

    async def connect_all(self, register: bool = True, quiet: bool = False) -> List[Any]:
        servers_config = self._load_config()
        all_tools = []
        # Connect to any configured servers but DO NOT auto-start vendored adapters.
        # Operators should install and start vendored MCPs manually (see third_party/).
        for name, config in servers_config.items():
            if not config.enabled:
                continue
            server = await self._connect_server(config)
            if server:
                self.servers[name] = server
                for tool_def in server.tools:
                    tool = create_mcp_tool(tool_def, server, self)
                    if register:
                        # Register tool into global registry so it appears in `tools list`
                        try:
                            register_tool_instance(tool)
                        except Exception:
                            pass
                    all_tools.append(tool)
                if not quiet:
                    print(f"[MCP] Connected to {name} with {len(server.tools)} tools")
        return all_tools

    async def connect_server(self, name: str, register: bool = True, quiet: bool = False) -> Optional[MCPServer]:
        servers_config = self._load_config()
        if name not in servers_config:
            return None
        config = servers_config[name]
        # Auto-start of vendored adapters (HexStrike/MetasploitMCP) has been
        # removed. Operators should install and run any third-party MCP adapters
        # manually under `third_party/` and configure `mcp_servers.json`.

        server = await self._connect_server(config)
        if server:
            self.servers[name] = server
            # Register tools for this server
            try:
                for tool_def in server.tools:
                    tool = create_mcp_tool(tool_def, server, self)
                    if register:
                        try:
                            register_tool_instance(tool)
                        except Exception:
                            pass
            except Exception:
                pass
            if not quiet:
                try:
                    print(f"[MCP] Connected to {name} with {len(server.tools)} tools")
                except Exception:
                    pass
        return server

    async def _connect_server(self, config: MCPServerConfig) -> Optional[MCPServer]:
        transport = None
        try:
            env = {**os.environ, **config.env}

            # Decide transport type:
            # - If args contain a http/sse transport or a --server http:// URL, use SSETransport
            # - Otherwise default to StdioTransport (spawn process and use stdio JSON-RPC)
            use_http = False
            http_url = None
            args_joined = " ".join(config.args or [])
            if "--transport http" in args_joined or "--transport sse" in args_joined:
                # Try to extract host/port from args
                try:
                    # naive parsing: look for --host <host> and --port <port>
                    host = None
                    port = None
                    for i, a in enumerate(config.args or []):
                        if a == "--host" and i + 1 < len(config.args):
                            host = config.args[i + 1]
                        if a == "--port" and i + 1 < len(config.args):
                            port = config.args[i + 1]
                    if host and port:
                        http_url = f"http://{host}:{port}/sse"
                except Exception:
                    http_url = None
                use_http = True
            # If args specify a --server URL, prefer that
            if not http_url:
                from urllib.parse import urlparse

                for i, a in enumerate(config.args or []):
                    if a == "--server" and i + 1 < len(config.args):
                        candidate = config.args[i + 1]
                        if isinstance(candidate, str) and candidate.startswith("http"):
                            # If the provided server URL doesn't include a path, default to the MCP SSE path
                            p = urlparse(candidate)
                            if p.path and p.path != "/":
                                http_url = candidate
                            else:
                                http_url = candidate.rstrip("/") + "/sse"
                            use_http = True
                            break

            if use_http and http_url:
                from .transport import SSETransport

                transport = SSETransport(url=http_url)
                await transport.connect()
            else:
                transport = StdioTransport(
                    command=config.command, args=config.args, env=env
                )
                await transport.connect()

            await transport.send(
                {
                    "jsonrpc": "2.0",
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {},
                        "clientInfo": {"name": "pentestagent", "version": "0.2.0"},
                    },
                    "id": self._get_next_id(),
                }
            )
            await transport.send(
                {"jsonrpc": "2.0", "method": "notifications/initialized"}
            )

            tools_response = await transport.send(
                {"jsonrpc": "2.0", "method": "tools/list", "id": self._get_next_id()}
            )
            tools = tools_response.get("result", {}).get("tools", [])

            return MCPServer(
                name=config.name,
                config=config,
                transport=transport,
                tools=tools,
                connected=True,
            )
        except Exception as e:
            # Clean up transport on failure
            if transport:
                try:
                    await transport.disconnect()
                except Exception:
                    pass
            print(f"[MCP] Failed to connect to {config.name}: {e}")
            return None

    async def call_tool(self, server_name: str, tool_name: str, arguments: dict) -> Any:
        server = self.servers.get(server_name)
        if not server or not server.connected:
            raise ValueError(f"Server '{server_name}' not connected")

        # Serialize all communication with this server to prevent:
        # - Message ID collisions
        # - Transport write interleaving
        # - Response routing issues
        async with server._lock:
            # Use 5 minute timeout for tool calls (scans can take a while)
            response = await server.transport.send(
                {
                    "jsonrpc": "2.0",
                    "method": "tools/call",
                    "params": {"name": tool_name, "arguments": arguments},
                    "id": self._get_next_id(),
                },
                timeout=300.0,
            )
        if "error" in response:
            raise RuntimeError(f"MCP error: {response['error'].get('message')}")
        return response.get("result", {}).get("content", [])

    async def disconnect_server(self, name: str):
        server = self.servers.get(name)
        if server:
            await server.disconnect()
            # Unregister MCP tools provided by this server
            try:
                for tool_def in server.tools:
                    mname = f"mcp_{server.name}_{tool_def.get('name')}"
                    try:
                        unregister_tool(mname)
                    except Exception:
                        pass
            except Exception:
                pass
            del self.servers[name]
        # If we started an adapter for this server, stop it as well
        adapter = self._started_adapters.pop(name, None)
        if adapter:
            try:
                stop = getattr(adapter, "stop", None)
                if stop:
                    if asyncio.iscoroutinefunction(stop):
                        await stop()
                    else:
                        loop = asyncio.get_running_loop()
                        await loop.run_in_executor(None, stop)
            except Exception:
                pass

    async def disconnect_all(self):
        for server in list(self.servers.values()):
            await server.disconnect()
        self.servers.clear()

    async def reconnect_all(self) -> List[Any]:
        """Disconnect all servers and reconnect them.

        Useful after cancellation leaves servers in a bad state.
        """
        # Disconnect all
        await self.disconnect_all()

        # Reconnect all configured servers
        return await self.connect_all()

    def get_server(self, name: str) -> Optional[MCPServer]:
        return self.servers.get(name)

    def get_all_servers(self) -> List[MCPServer]:
        return list(self.servers.values())

    def is_connected(self, name: str) -> bool:
        server = self.servers.get(name)
        return server is not None and server.connected

    async def start_control_server(self, path: Optional[str] = None) -> str:
        """Start a lightweight UNIX-domain socket control server.

        The control server accepts newline-delimited JSON requests. Supported
        commands:
          {"cmd": "status"} -> returns connected servers and counts
          {"cmd": "list_tools"} -> returns list of MCP tools (name, server, description)

        Returns the path of the socket in use.
        """
        if not path:
            path = str(Path.home() / ".pentestagent" / "mcp.sock")

        sock_path = Path(path)
        # Ensure parent exists
        sock_path.parent.mkdir(parents=True, exist_ok=True)

        # Remove stale socket if present
        try:
            if sock_path.exists():
                try:
                    sock_path.unlink()
                except Exception:
                    pass
        except Exception:
            pass

        async def _handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
            try:
                data = await reader.readline()
                if not data:
                    return
                try:
                    req = json.loads(data.decode("utf-8"))
                except Exception:
                    resp = {"status": "error", "error": "invalid_json"}
                    writer.write((json.dumps(resp) + "\n").encode("utf-8"))
                    await writer.drain()
                    return

                cmd = req.get("cmd") if isinstance(req, dict) else None
                if cmd == "status":
                    servers = []
                    for name, s in self.servers.items():
                        servers.append({"name": name, "connected": bool(s.connected), "tool_count": len(s.tools)})
                    resp = {"status": "ok", "servers": servers}
                    writer.write((json.dumps(resp) + "\n").encode("utf-8"))
                    await writer.drain()
                elif cmd == "list_tools":
                    tools = []
                    for sname, s in self.servers.items():
                        for t in s.tools:
                            tools.append({"name": t.get("name"), "server": sname, "description": t.get("description", "")})
                    resp = {"status": "ok", "tools": tools}
                    writer.write((json.dumps(resp) + "\n").encode("utf-8"))
                    await writer.drain()
                elif cmd == "call_tool":
                    # Expecting: {"cmd":"call_tool","server":"name","tool":"tool_name","args":{...}}
                    server_name = req.get("server")
                    tool_name = req.get("tool")
                    arguments = req.get("args", {}) if isinstance(req.get("args", {}), dict) else {}
                    if not server_name or not tool_name:
                        writer.write((json.dumps({"status": "error", "error": "missing_parameters"}) + "\n").encode("utf-8"))
                        await writer.drain()
                        return
                    try:
                        # perform the tool call
                        result = await self.call_tool(server_name, tool_name, arguments)
                        writer.write((json.dumps({"status": "ok", "result": result}) + "\n").encode("utf-8"))
                        await writer.drain()
                    except Exception as e:
                        writer.write((json.dumps({"status": "error", "error": "call_failed", "message": str(e)}) + "\n").encode("utf-8"))
                        await writer.drain()
                    
                else:
                    resp = {"status": "error", "error": "unknown_cmd"}
                    writer.write((json.dumps(resp) + "\n").encode("utf-8"))
                    await writer.drain()
            except Exception:
                try:
                    writer.write((json.dumps({"status": "error", "error": "internal"}) + "\n").encode("utf-8"))
                    await writer.drain()
                except Exception:
                    pass
            finally:
                try:
                    writer.close()
                except Exception:
                    pass

        # Start the asyncio unix server
        loop = asyncio.get_running_loop()
        server = await asyncio.start_unix_server(_handle, path=path)
        self._control_server = server
        self._control_path = Path(path)

        # Keep server serving in background task
        self._control_task = loop.create_task(server.serve_forever())
        # Restrict socket access to current user where possible
        try:
            os.chmod(path, 0o600)
        except Exception:
            pass
        return path

    async def stop_control_server(self):
        try:
            if self._control_task:
                self._control_task.cancel()
                self._control_task = None
            if self._control_server:
                self._control_server.close()
                try:
                    await self._control_server.wait_closed()
                except Exception:
                    pass
                self._control_server = None
            if self._control_path and self._control_path.exists():
                try:
                    self._control_path.unlink()
                except Exception:
                    pass
            self._control_path = None
        except Exception:
            pass
