#!/usr/bin/env python3
"""Subprocess lifecycle manager for external tools.

Manages bettercap, tcpdump, Responder, mitmproxy, ntlmrelayx as
supervised subprocesses with:
  - Start/stop/restart with configurable retry (max 3, exponential backoff)
  - stdout/stderr capture
  - PID tracking for kill_switch integration
  - Resource monitoring via /proc/PID/stat
  - Process disguise integration (prctl rename)
  - Health check callbacks
  - Graceful shutdown (SIGTERM -> 5s -> SIGKILL)
  - Crash callbacks (e.g., ARP restoration on bettercap crash)
"""

import os
import time
import signal
import logging
import subprocess
import threading
from dataclasses import dataclass, field
from typing import Callable, Optional

from core.bus import EventBus

logger = logging.getLogger(__name__)


@dataclass
class ManagedTool:
    """Configuration and runtime state for a managed subprocess."""
    name: str
    binary: str
    args: list = field(default_factory=list)
    health_check: Optional[Callable] = None
    crash_callback: Optional[Callable] = None
    max_restarts: int = 3
    disguise_name: str = ""
    env: dict = field(default_factory=dict)
    cwd: Optional[str] = None

    # Runtime state
    process: Optional[subprocess.Popen] = field(default=None, repr=False)
    pid: Optional[int] = None
    restart_count: int = 0
    started_at: Optional[float] = None
    _stdout_thread: Optional[threading.Thread] = field(default=None, repr=False)
    _stderr_thread: Optional[threading.Thread] = field(default=None, repr=False)
    stdout_lines: list = field(default_factory=list, repr=False)
    stderr_lines: list = field(default_factory=list, repr=False)
    _max_log_lines: int = 1000


class ToolManager:
    """Manages external tool subprocesses with supervision.

    Usage:
        tm = ToolManager(bus)
        tm.register(name="bettercap", binary="/usr/local/bin/bettercap",
                     args=["-api-rest-address", "127.0.0.1"],
                     health_check=lambda: api.is_alive(),
                     crash_callback=on_crash, max_restarts=3,
                     disguise_name="networkd-dispatcher")
        tm.start("bettercap")
        tm.stop("bettercap")
    """

    SIGTERM_TIMEOUT = 5.0  # seconds to wait after SIGTERM before SIGKILL

    def __init__(self, bus: EventBus):
        self.bus = bus
        self._tools: dict[str, ManagedTool] = {}
        self._lock = threading.Lock()
        self._monitor_thread: Optional[threading.Thread] = None
        self._running = False

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def start_monitoring(self) -> None:
        """Start the background health/crash monitor."""
        if self._running:
            return
        self._running = True
        self._monitor_thread = threading.Thread(
            target=self._monitor_loop, daemon=True, name="sensor-tool-monitor"
        )
        self._monitor_thread.start()

    def stop_monitoring(self) -> None:
        """Stop monitoring thread."""
        self._running = False
        if self._monitor_thread and self._monitor_thread.is_alive():
            self._monitor_thread.join(timeout=3.0)

    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------

    def register(self, name: str, binary: str, args: list = None,
                 health_check: Callable = None, crash_callback: Callable = None,
                 max_restarts: int = 3, disguise_name: str = "",
                 env: dict = None, cwd: str = None) -> None:
        """Register a tool for management."""
        with self._lock:
            self._tools[name] = ManagedTool(
                name=name,
                binary=binary,
                args=args or [],
                health_check=health_check,
                crash_callback=crash_callback,
                max_restarts=max_restarts,
                disguise_name=disguise_name,
                env=env or {},
                cwd=cwd,
            )
        logger.info("Registered tool: %s (%s)", name, binary)

    # ------------------------------------------------------------------
    # Start / Stop / Restart
    # ------------------------------------------------------------------

    def start(self, name: str) -> bool:
        """Start a registered tool subprocess."""
        with self._lock:
            tool = self._tools.get(name)
            if tool is None:
                logger.error("Tool %s not registered", name)
                return False

            if tool.process and tool.process.poll() is None:
                logger.warning("Tool %s already running (pid=%d)", name, tool.pid)
                return True

            return self._start_tool(tool)

    def _start_tool(self, tool: ManagedTool) -> bool:
        """Internal: launch the subprocess. Caller must hold _lock."""
        cmd = [tool.binary] + tool.args
        env = dict(os.environ)
        env.update(tool.env)

        try:
            proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=env,
                cwd=tool.cwd,
                preexec_fn=self._make_preexec(tool.disguise_name) if tool.disguise_name else None,
            )
        except FileNotFoundError:
            logger.error("Binary not found: %s", tool.binary)
            return False
        except Exception as e:
            logger.error("Failed to start %s: %s", tool.name, e)
            return False

        tool.process = proc
        tool.pid = proc.pid
        tool.started_at = time.time()
        tool.stdout_lines.clear()
        tool.stderr_lines.clear()

        # Stdout/stderr reader threads
        tool._stdout_thread = threading.Thread(
            target=self._pipe_reader, args=(proc.stdout, tool.stdout_lines, tool._max_log_lines),
            daemon=True, name=f"bb-{tool.name}-stdout",
        )
        tool._stderr_thread = threading.Thread(
            target=self._pipe_reader, args=(proc.stderr, tool.stderr_lines, tool._max_log_lines),
            daemon=True, name=f"bb-{tool.name}-stderr",
        )
        tool._stdout_thread.start()
        tool._stderr_thread.start()

        self.bus.emit("TOOL_RESTARTED" if tool.restart_count > 0 else "MODULE_STARTED",
                      {"tool": tool.name, "pid": proc.pid, "restart_count": tool.restart_count},
                      source_module="tool_manager")
        logger.info("Started tool %s (pid=%d)", tool.name, proc.pid)
        return True

    def stop(self, name: str) -> bool:
        """Stop a tool. SIGTERM -> 5s timeout -> SIGKILL."""
        with self._lock:
            tool = self._tools.get(name)
            if tool is None:
                return True
            return self._stop_tool(tool)

    def _stop_tool(self, tool: ManagedTool) -> bool:
        """Internal: stop subprocess. Caller must hold _lock."""
        proc = tool.process
        if proc is None or proc.poll() is not None:
            tool.process = None
            tool.pid = None
            return True

        # Graceful: SIGTERM
        try:
            proc.terminate()
        except OSError:
            pass

        try:
            proc.wait(timeout=self.SIGTERM_TIMEOUT)
        except subprocess.TimeoutExpired:
            # Forceful: SIGKILL
            logger.warning("Tool %s did not stop after SIGTERM, sending SIGKILL", tool.name)
            try:
                proc.kill()
                proc.wait(timeout=2.0)
            except Exception:
                pass

        tool.process = None
        tool.pid = None
        tool.started_at = None
        logger.info("Stopped tool %s", tool.name)
        return True

    def restart(self, name: str) -> bool:
        """Stop then start a tool, resetting restart count."""
        with self._lock:
            tool = self._tools.get(name)
            if tool is None:
                return False
            self._stop_tool(tool)
            tool.restart_count = 0
            return self._start_tool(tool)

    def stop_all(self) -> None:
        """Stop all managed tools."""
        with self._lock:
            for tool in self._tools.values():
                self._stop_tool(tool)
        self.stop_monitoring()

    # ------------------------------------------------------------------
    # Status / resources
    # ------------------------------------------------------------------

    def get_status(self, name: str) -> dict:
        """Get tool status including resource usage from /proc."""
        with self._lock:
            tool = self._tools.get(name)
            if tool is None:
                return {"error": "not registered"}

        alive = tool.process is not None and tool.process.poll() is None
        result = {
            "name": tool.name,
            "running": alive,
            "pid": tool.pid,
            "uptime": time.time() - tool.started_at if tool.started_at and alive else 0,
            "restart_count": tool.restart_count,
        }

        if alive and tool.pid:
            resources = self._read_proc_stats(tool.pid)
            result.update(resources)

        return result

    def get_all_status(self) -> dict:
        """Get status for all registered tools."""
        with self._lock:
            names = list(self._tools.keys())
        return {n: self.get_status(n) for n in names}

    def get_all_pids(self) -> list[int]:
        """Return PIDs of all running tools (for kill_switch)."""
        pids = []
        with self._lock:
            for tool in self._tools.values():
                if tool.pid and tool.process and tool.process.poll() is None:
                    pids.append(tool.pid)
        return pids

    @staticmethod
    def _read_proc_stats(pid: int) -> dict:
        """Read RAM/CPU from /proc/PID/stat and /proc/PID/statm."""
        result = {"ram_mb": 0.0, "cpu_pct": 0.0}
        try:
            with open(f"/proc/{pid}/statm", "r") as f:
                parts = f.read().split()
                # statm: size resident shared text lib data dt (in pages)
                page_size = os.sysconf("SC_PAGE_SIZE")
                result["ram_mb"] = round(int(parts[1]) * page_size / (1024 * 1024), 1)
        except (FileNotFoundError, PermissionError, IndexError, ValueError):
            pass

        try:
            with open(f"/proc/{pid}/stat", "r") as f:
                parts = f.read().split()
                # Fields 14,15 = utime, stime in clock ticks
                utime = int(parts[13])
                stime = int(parts[14])
                total_ticks = utime + stime
                hz = os.sysconf("SC_CLK_TCK")
                uptime_f = open("/proc/uptime", "r")
                system_uptime = float(uptime_f.read().split()[0])
                uptime_f.close()
                start_ticks = int(parts[21])
                proc_uptime = system_uptime - (start_ticks / hz)
                if proc_uptime > 0:
                    result["cpu_pct"] = round((total_ticks / hz) / proc_uptime * 100, 1)
        except (FileNotFoundError, PermissionError, IndexError, ValueError, ZeroDivisionError):
            pass

        return result

    # ------------------------------------------------------------------
    # Monitor loop
    # ------------------------------------------------------------------

    def _monitor_loop(self) -> None:
        """Background loop: detect crashes, run health checks, auto-restart."""
        while self._running:
            time.sleep(5.0)
            with self._lock:
                tools = list(self._tools.values())

            for tool in tools:
                if tool.process is None:
                    continue

                rc = tool.process.poll()
                if rc is not None:
                    # Process exited
                    logger.warning("Tool %s exited (rc=%s)", tool.name, rc)
                    self.bus.emit("TOOL_CRASHED",
                                 {"tool": tool.name, "exit_code": rc,
                                  "restart_count": tool.restart_count},
                                 source_module="tool_manager")

                    # Call crash callback OUTSIDE lock to prevent deadlock
                    callback = tool.crash_callback
                    restart_needed = tool.restart_count < tool.max_restarts
                    
                    # Release lock before callback
                    if callback:
                        try:
                            callback()
                        except Exception:
                            logger.exception("Crash callback failed for %s", tool.name)
                    
                    # Re-acquire lock for restart decision
                    with self._lock:
                        # Auto-restart with exponential backoff
                        if tool.restart_count < tool.max_restarts:
                            backoff = min(2 ** tool.restart_count, 30)
                            logger.info("Restarting %s in %ds (attempt %d/%d)",
                                        tool.name, backoff, tool.restart_count + 1,
                                        tool.max_restarts)
                            time.sleep(backoff)
                            tool.restart_count += 1
                            self._start_tool(tool)
                        else:
                            logger.error("Tool %s exceeded max restarts (%d), giving up",
                                         tool.name, tool.max_restarts)
                            tool.process = None
                            tool.pid = None

                elif tool.health_check:
                    # Process alive — run health check
                    try:
                        if not tool.health_check():
                            logger.warning("Health check failed for %s", tool.name)
                    except Exception:
                        logger.exception("Health check error for %s", tool.name)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _pipe_reader(pipe, line_buffer: list, max_lines: int) -> None:
        """Read lines from a subprocess pipe into a bounded buffer."""
        try:
            for line in iter(pipe.readline, b""):
                try:
                    decoded = line.decode("utf-8", errors="replace").rstrip()
                except Exception:
                    decoded = repr(line)
                line_buffer.append(decoded)
                if len(line_buffer) > max_lines:
                    line_buffer.pop(0)
        except Exception:
            pass
        finally:
            try:
                pipe.close()
            except Exception:
                pass

    @staticmethod
    def _make_preexec(disguise_name: str):
        """Return a preexec_fn that renames the process via prctl."""
        def _preexec():
            try:
                import ctypes
                libc = ctypes.CDLL("libc.so.6", use_errno=True)
                PR_SET_NAME = 15
                name_bytes = disguise_name[:15].encode("utf-8")
                libc.prctl(PR_SET_NAME, name_bytes, 0, 0, 0)
            except Exception:
                pass
        return _preexec

    def get_output(self, name: str, stream: str = "stderr",
                   lines: int = 50) -> list[str]:
        """Get recent stdout/stderr lines from a tool."""
        with self._lock:
            tool = self._tools.get(name)
            if tool is None:
                return []
        buf = tool.stderr_lines if stream == "stderr" else tool.stdout_lines
        return list(buf[-lines:])
