"""Adapter to manage a vendored Metasploit MCP server.

This follows the same lightweight pattern as the Hexstrike adapter: it
expects the MetasploitMCP repository to be vendored under
``third_party/MetasploitMCP`` (or a custom path provided by the caller).
The adapter starts the server as a background subprocess and performs a
health check on a configurable port.
"""

import asyncio
import os
import shutil
from pathlib import Path
from typing import Optional
import time
import signal

try:
    import aiohttp
except Exception:
    aiohttp = None


LOOT_DIR = Path("loot/artifacts")
LOOT_DIR.mkdir(parents=True, exist_ok=True)
LOG_FILE = LOOT_DIR / "metasploit_mcp.log"


class MetasploitAdapter:
    """Manage a vendored Metasploit MCP server under `third_party/MetasploitMCP`.

    Usage:
        adapter = MetasploitAdapter()
        await adapter.start()
        # ... use MCPManager to connect to the server
        await adapter.stop()
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 7777,
        python_cmd: str = "python3",
        server_script: Optional[Path] = None,
        cwd: Optional[Path] = None,
        env: Optional[dict] = None,
        transport: str = "http",
    ) -> None:
        self.host = host
        self.port = int(port)
        self.python_cmd = python_cmd
        # Vendored project uses 'MetasploitMCP.py' as the main entrypoint
        self.server_script = (
            server_script or Path("third_party/MetasploitMCP/MetasploitMCP.py")
        )
        self.cwd = cwd or Path.cwd()
        self.env = {**os.environ, **(env or {})}
        self.transport = transport

        self._process: Optional[asyncio.subprocess.Process] = None
        self._reader_task: Optional[asyncio.Task] = None
        self._msfrpcd_proc: Optional[asyncio.subprocess.Process] = None

    def _build_command(self):
        cmd = [self.python_cmd, str(self.server_script)]
        # Prefer explicit transport when starting vendored server from adapter
        if self.transport:
            cmd += ["--transport", str(self.transport)]
        # When running HTTP, ensure host/port are provided
        if str(self.transport).lower() in ("http", "sse"):
            cmd += ["--host", str(self.host), "--port", str(self.port)]
        else:
            # For other transports, allow default args
            cmd += ["--port", str(self.port)]
        return cmd

    async def _start_msfrpcd_if_needed(self) -> None:
        """Start `msfrpcd` if it's not already reachable at MSF_SERVER:MSF_PORT.

        This starts `msfrpcd` as a child process (no sudo) using MSF_* env
        values if available. It's intentionally conservative: if the RPC
        endpoint is already listening we won't try to start a new daemon.
        """
        try:
            msf_server = str(self.env.get("MSF_SERVER", "127.0.0.1"))
            msf_port = int(self.env.get("MSF_PORT", 55553))
        except Exception:
            msf_server = "127.0.0.1"
            msf_port = 55553

        # Quick socket check to see if msfrpcd is already listening
        import socket

        try:
            with socket.create_connection((msf_server, msf_port), timeout=1):
                return
        except Exception:
            pass

        # If msfrpcd not available on path, skip starting
        if not shutil.which("msfrpcd"):
            return

        msf_user = str(self.env.get("MSF_USER", "msf"))
        msf_password = str(self.env.get("MSF_PASSWORD", ""))
        msf_ssl = str(self.env.get("MSF_SSL", "false")).lower() in ("1", "true", "yes", "y")

        # Build args for msfrpcd (no sudo). Use -S (SSL optional) flag only if requested.
        args = ["msfrpcd", "-U", msf_user, "-P", msf_password, "-a", msf_server, "-p", str(msf_port)]
        if msf_ssl:
            args.append("-S")

        try:
            resolved = shutil.which("msfrpcd") or "msfrpcd"
            self._msfrpcd_proc = await asyncio.create_subprocess_exec(
                resolved,
                *args[1:],
                cwd=str(self.cwd),
                env=self.env,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                start_new_session=True,
            )
            # Start reader to capture msfrpcd logs
            loop = asyncio.get_running_loop()
            loop.create_task(self._capture_msfrpcd_output())
            # Give msfrpcd a moment to start
            await asyncio.sleep(0.5)
        except Exception:
            return

    async def _capture_msfrpcd_output(self) -> None:
        if not self._msfrpcd_proc or not self._msfrpcd_proc.stdout:
            return
        try:
            with LOG_FILE.open("ab") as fh:
                while True:
                    line = await self._msfrpcd_proc.stdout.readline()
                    if not line:
                        break
                    fh.write(b"[msfrpcd] " + line)
                    fh.flush()
        except asyncio.CancelledError:
            pass
        except Exception:
            pass

    async def start(self, background: bool = True, timeout: int = 30) -> bool:
        """Start the vendored Metasploit MCP server.

        Returns True if the server started and passed health check within
        `timeout` seconds.
        """
        if not self.server_script.exists():
            raise FileNotFoundError(
                f"Metasploit MCP server script not found at {self.server_script}."
            )

        if self._process and self._process.returncode is None:
            return await self.health_check(timeout=1)

        # If running in HTTP/SSE mode, ensure msfrpcd is started and reachable
        if str(self.transport).lower() in ("http", "sse"):
            try:
                await self._start_msfrpcd_if_needed()
            except Exception:
                pass

        cmd = self._build_command()
        resolved = shutil.which(self.python_cmd) or self.python_cmd

        self._process = await asyncio.create_subprocess_exec(
            resolved,
            *cmd[1:],
            cwd=str(self.cwd),
            env=self.env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            start_new_session=True,
        )

        # Log PID
        try:
            pid = getattr(self._process, "pid", None)
            if pid:
                with LOG_FILE.open("a") as fh:
                    fh.write(f"[MetasploitAdapter] started pid={pid}\n")
        except Exception:
            pass

        # Start background reader
        loop = asyncio.get_running_loop()
        self._reader_task = loop.create_task(self._capture_output())

        try:
            return await self.health_check(timeout=timeout)
        except Exception:
            return False

    async def _capture_output(self) -> None:
        if not self._process or not self._process.stdout:
            return

        try:
            with LOG_FILE.open("ab") as fh:
                while True:
                    line = await self._process.stdout.readline()
                    if not line:
                        break
                    fh.write(line)
                    fh.flush()
        except asyncio.CancelledError:
            pass
        except Exception:
            pass

    async def stop(self, timeout: int = 5) -> None:
        proc = self._process
        if not proc:
            return

        try:
            proc.terminate()
            await asyncio.wait_for(proc.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            try:
                proc.kill()
            except Exception:
                pass
        except Exception:
            pass

        self._process = None

        if self._reader_task and not self._reader_task.done():
            self._reader_task.cancel()
            try:
                await self._reader_task
            except Exception:
                pass

        # Stop msfrpcd if we started it
        try:
            msf_proc = self._msfrpcd_proc
            if msf_proc:
                try:
                    msf_proc.terminate()
                    await asyncio.wait_for(msf_proc.wait(), timeout=timeout)
                except asyncio.TimeoutError:
                    try:
                        msf_proc.kill()
                    except Exception:
                        pass
        except Exception:
            pass
        finally:
            self._msfrpcd_proc = None

    def stop_sync(self, timeout: int = 5) -> None:
        proc = self._process
        if not proc:
            return

        try:
            pid = getattr(proc, "pid", None)
            if pid:
                try:
                    pgid = os.getpgid(pid)
                    os.killpg(pgid, signal.SIGTERM)
                except Exception:
                    try:
                        os.kill(pid, signal.SIGTERM)
                    except Exception:
                        pass

                end = time.time() + float(timeout)
                while time.time() < end:
                    ret = getattr(proc, "returncode", None)
                    if ret is not None:
                        break
                    time.sleep(0.1)

                try:
                    pgid = os.getpgid(pid)
                    os.killpg(pgid, signal.SIGKILL)
                except Exception:
                    try:
                        os.kill(pid, signal.SIGKILL)
                    except Exception:
                        pass
        except Exception:
            pass

    def __del__(self):
        try:
            self.stop_sync()
        except Exception:
            pass
        try:
            self._process = None
        except Exception:
            pass

    async def health_check(self, timeout: int = 5) -> bool:
        url = f"http://{self.host}:{self.port}/health"

        if aiohttp:
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.get(url, timeout=timeout) as resp:
                        return resp.status == 200
            except Exception:
                return False

        import urllib.request

        def _check():
            try:
                with urllib.request.urlopen(url, timeout=timeout) as r:
                    return r.status == 200
            except Exception:
                return False

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, _check)

    def is_running(self) -> bool:
        return self._process is not None and self._process.returncode is None


__all__ = ["MetasploitAdapter"]
