#!/usr/bin/env python3
"""Module lifecycle manager.

Each SystemMonitor module runs in its own multiprocessing.Process.
The Engine handles:
  - Loading module classes from the modules/ tree
  - Dependency resolution (topological sort)
  - Start / stop / restart / start_all / stop_all
  - Resource budgeting per hardware tier
  - Interface conflict detection
  - Scope enforcement
  - Engagement phase gating (passive_only for N days, then active_allowed)
"""

import os
import time
import signal
import logging
import importlib
import multiprocessing
from collections import defaultdict, deque
from typing import Optional, Type

from core.bus import EventBus, Event
from core.state import StateManager
from modules.base import BaseModule
from utils.networking import detect_interface_with_retry

logger = logging.getLogger(__name__)

# Hardware tier definitions — module limits and resource ceilings
HARDWARE_TIERS = {
    "opi_zero3": {
        "max_passive": 16,
        "max_active": 6,
        "max_ram_mb": 3072,
        "max_cpu_pct": 80,
        "allow_mitmproxy": True,
        "allow_bettercap": True,
    },
    "pi3b": {
        "max_passive": 8,
        "max_active": 2,
        "max_stealth": 3,
        "max_other": 4,
        "max_ram_mb": 700,  # 900MB total - ~200MB OS overhead
        "max_cpu_pct": 75,
        "allow_mitmproxy": False,
        "allow_bettercap": False,
    },
    "pi_zero": {
        "max_passive": 4,
        "max_active": 0,
        "max_ram_mb": 400,
        "max_cpu_pct": 70,
        "allow_mitmproxy": False,
        "allow_bettercap": False,
    },
    "generic": {
        "max_passive": 16,
        "max_active": 99,
        "max_ram_mb": 0,  # 0 = no limit
        "max_cpu_pct": 90,
        "allow_mitmproxy": True,
        "allow_bettercap": True,
    },
}


def detect_hardware_tier() -> str:
    """Detect hardware tier from /proc/device-tree/model and /proc/cpuinfo."""
    model = ""
    try:
        with open("/proc/device-tree/model", "r") as f:
            model = f.read().strip().lower()
    except (FileNotFoundError, PermissionError):
        pass

    if "orange pi zero3" in model or "orange pi zero 3" in model:
        return "opi_zero3"

    # Check cpuinfo for known SoCs
    try:
        with open("/proc/cpuinfo", "r") as f:
            cpuinfo = f.read().lower()
    except (FileNotFoundError, PermissionError):
        cpuinfo = ""

    if "allwinner" in cpuinfo and "h618" in cpuinfo:
        return "opi_zero3"

    if "raspberry pi zero 2" in model:
        return "pi_zero"

    if "raspberry pi 3" in model:
        return "pi3b"

    # BCM2710/BCM2837 without a clear model string — Pi 3B family
    if "bcm2837" in cpuinfo or ("bcm2710" in cpuinfo and "raspberry pi zero 2" not in model):
        return "pi3b"

    return "generic"


def _topo_sort(modules: dict[str, "ModuleEntry"]) -> list[str]:
    """Topological sort of modules by dependency. Returns ordered list of names."""
    in_degree: dict[str, int] = defaultdict(int)
    graph: dict[str, list[str]] = defaultdict(list)
    all_names = set(modules.keys())

    for name, entry in modules.items():
        for dep in entry.module_class.dependencies:
            if dep in all_names:
                graph[dep].append(name)
                in_degree[name] += 1
        if name not in in_degree:
            in_degree[name] = 0

    queue = deque(n for n, d in in_degree.items() if d == 0)
    order = []
    while queue:
        node = queue.popleft()
        order.append(node)
        for neighbor in graph[node]:
            in_degree[neighbor] -= 1
            if in_degree[neighbor] == 0:
                queue.append(neighbor)

    if len(order) != len(all_names):
        missing = all_names - set(order)
        logger.error("Circular dependency detected involving: %s", missing)
        # Append remaining modules anyway (best-effort)
        order.extend(missing)

    return order


class ModuleEntry:
    """Registry entry for a loaded module."""
    __slots__ = ("module_class", "config", "process", "started_at")

    def __init__(self, module_class: Type[BaseModule], config: dict):
        self.module_class = module_class
        self.config = config
        self.process: Optional[multiprocessing.Process] = None
        self.started_at: Optional[float] = None


def _module_runner(module_class: Type[BaseModule], bus_queue: multiprocessing.Queue,
                   state_db_path: str, config: dict, engine_pipe) -> None:
    """Entry point for each module process.

    Instantiates the module, wires up a local EventBus that forwards
    events through the shared mp.Queue, and calls module.start().
    """
    # Build a lightweight bus proxy that pushes to the shared queue
    class _BusProxy:
        def __init__(self, queue):
            self._q = queue

        def publish(self, event):
            try:
                self._q.put_nowait(event)
            except Exception:
                pass

        def emit(self, event_type, payload, source_module=""):
            self.publish(Event(
                event_type=event_type,
                payload=payload,
                source_module=source_module,
            ))

        def subscribe(self, *a, **kw):
            pass  # Subscriptions only work in the main process

    bus_proxy = _BusProxy(bus_queue)
    state = StateManager(db_path=state_db_path)
    state.start()

    # CaptureBus reader thread is not inherited across fork — restart it in
    # this subprocess so the module's subscribe() calls reach a live reader.
    for key in ("capture_bus", "_capture_bus"):
        cb = config.get(key)
        if cb is not None:
            cb._running = False
            cb._capture_thread = None
            cb._sock = None
            cb._subscribers = []
            cb.start()
            break

    module = module_class(bus=bus_proxy, state=state, config=config)

    # Signal handlers for graceful shutdown
    def _handle_term(sig, frame):
        try:
            module.stop()
        except Exception:
            pass
        state.stop()
        os._exit(0)

    signal.signal(signal.SIGTERM, _handle_term)
    signal.signal(signal.SIGINT, _handle_term)

    try:
        module.start()
        # start() is non-blocking — keep subprocess alive while module threads run
        while module._running:
            time.sleep(1)
    except Exception:
        logger.exception("Module %s crashed in start()", module.name)
        bus_proxy.emit("MODULE_ERROR", {"module": module.name, "error": "start_crash"},
                       source_module=module.name)
    finally:
        try:
            module.stop()
        except Exception:
            pass
        state.stop()


class Engine:
    """Module lifecycle manager.

    Manages multiprocess module instances with dependency resolution,
    hardware tier enforcement, scope checking, and phase gating.
    """

    def __init__(self, bus: EventBus, state: StateManager, config: dict):
        self.bus = bus
        self.state = state
        self.config = config
        self._modules: dict[str, ModuleEntry] = {}
        self._tier = config.get("device", {}).get("platform", "auto")
        if self._tier == "auto":
            self._tier = detect_hardware_tier()
        self._tier_limits = HARDWARE_TIERS.get(self._tier, HARDWARE_TIERS["generic"])
        self._scope = config.get("scope", {})
        self._phase = config.get("engagement_phase", {}).get("mode", "passive_only")

        # Resolve "auto" interface to actual interface name with retry fallback
        if self.config.get("network", {}).get("primary_interface") == "auto":
            actual_iface = detect_interface_with_retry(
                max_retries=3,
                retry_delay=5,
                exponential=True,
                config_interface=None
            )
            if actual_iface:
                self.config.setdefault("network", {})["primary_interface"] = actual_iface
                logger.info("Resolved auto interface to: %s (with retry)", actual_iface)
        self.capture_bus = None
        self._lock = multiprocessing.Lock()
        logger.info("Engine initialized (tier=%s, phase=%s)", self._tier, self._phase)

    # ------------------------------------------------------------------
    # Module registration
    # ------------------------------------------------------------------

    def register(self, module_class: Type[BaseModule], config: dict = None) -> None:
        """Register a module class for management."""
        name = module_class.name
        if name in self._modules:
            logger.warning("Module %s already registered, replacing", name)
        self._modules[name] = ModuleEntry(module_class, config or {})
        logger.debug("Registered module: %s (type=%s, priority=%d)",
                     name, module_class.module_type, module_class.priority)

    def load_module(self, module_path: str, config: dict = None) -> None:
        """Dynamically load a module from a dotted Python path.

        e.g. "modules.passive.dns_logger.DnsLogger"
        """
        parts = module_path.rsplit(".", 1)
        if len(parts) != 2:
            raise ValueError(f"Expected 'package.ClassName', got: {module_path}")
        mod = importlib.import_module(parts[0])
        cls = getattr(mod, parts[1])
        if not (isinstance(cls, type) and issubclass(cls, BaseModule)):
            raise TypeError(f"{module_path} is not a BaseModule subclass")
        self.register(cls, config)

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

    def start(self, name: str) -> bool:
        """Start a single module by name. Returns True on success."""
        with self._lock:
            entry = self._modules.get(name)
            if entry is None:
                logger.error("Module %s not registered", name)
                return False

            if entry.process and entry.process.is_alive():
                logger.warning("Module %s already running", name)
                return True

            # Phase gating
            if not self._phase_allows(entry.module_class):
                logger.warning("Module %s blocked by phase=%s", name, self._phase)
                return False

            # Tier limits
            if not self._tier_allows(entry.module_class):
                logger.warning("Module %s blocked by tier=%s limits", name, self._tier)
                return False

            # Dependency check
            for dep in entry.module_class.dependencies:
                dep_entry = self._modules.get(dep)
                if dep_entry is None or not (dep_entry.process and dep_entry.process.is_alive()):
                    logger.error("Module %s requires %s which is not running", name, dep)
                    return False

            proc = multiprocessing.Process(
                target=_module_runner,
                args=(
                    entry.module_class,
                    self.bus.get_queue(),
                    self.state._db_path,
                    entry.config,
                    None,
                ),
                name=f"bb-{name}",
                daemon=True,
            )
            proc.start()
            entry.process = proc
            entry.started_at = time.time()

            self.state.set_module_status(name, "running", pid=proc.pid)
            self.bus.emit("MODULE_STARTED", {"module": name, "pid": proc.pid},
                          source_module="engine")
            logger.info("Started module %s (pid=%d)", name, proc.pid)
            return True

    def stop(self, name: str, timeout: float = 5.0) -> bool:
        """Stop a module. SIGTERM -> wait -> SIGKILL."""
        with self._lock:
            entry = self._modules.get(name)
            if entry is None or entry.process is None:
                return True

            proc = entry.process
            if proc.is_alive():
                proc.terminate()
                proc.join(timeout=timeout)
                if proc.is_alive():
                    logger.warning("Module %s did not stop, sending SIGKILL", name)
                    proc.kill()
                    proc.join(timeout=2.0)

            entry.process = None
            entry.started_at = None
            self.state.set_module_status(name, "stopped")
            self.bus.emit("MODULE_STOPPED", {"module": name}, source_module="engine")
            logger.info("Stopped module %s", name)
            return True

    def restart(self, name: str) -> bool:
        """Stop then start a module."""
        self.stop(name)
        time.sleep(0.2)
        return self.start(name)

    def start_all(self) -> dict[str, bool]:
        """Start all registered modules in dependency order."""
        # Clear stale PIDs from any previous run so the watchdog does not
        # report modules from the old session as dead.
        self.state.reset_module_status()

        order = _topo_sort(self._modules)

        # Instantiate and start CaptureBus for passive modules
        iface = self.config.get("network", {}).get("primary_interface")
        if iface:
            try:
                from core.capture_bus import CaptureBus
                self.capture_bus = CaptureBus(interface=iface)
                self.capture_bus.start()
                logger.info("CaptureBus started on %s", iface)
            except Exception as e:
                logger.error("Failed to start CaptureBus: %s", e)
                self.capture_bus = None

        # Inject capture_bus into module configs for passive modules
        # Some modules use "capture_bus", others "_capture_bus" — inject both
        for name, entry in self._modules.items():
            if getattr(entry.module_class, "requires_capture_bus", False):
                entry.config["capture_bus"] = self.capture_bus
                entry.config["_capture_bus"] = self.capture_bus

        results = {}
        for name in order:
            results[name] = self.start(name)

        # Post-startup liveness check: wait for module processes to settle,
        # then verify each one is still alive and log a clear summary.
        time.sleep(3)
        alive, dead = [], []
        with self._lock:
            for name, entry in self._modules.items():
                if results.get(name):
                    if entry.process and entry.process.is_alive():
                        alive.append(name)
                    else:
                        dead.append(name)
                        # Update state so watchdog doesn't double-report
                        self.state.set_module_status(name, "stopped")

        if dead:
            logger.error(
                "Startup check: %d modules died immediately after launch: %s",
                len(dead), ", ".join(dead),
            )
        logger.info(
            "Startup check complete — running: %d, failed: %d",
            len(alive), len(dead),
        )

        return results

    def stop_all(self, timeout: float = 5.0) -> None:
        """Stop all running modules in reverse dependency order."""
        order = _topo_sort(self._modules)
        for name in reversed(order):
            self.stop(name, timeout=timeout)
        
        # Stop CaptureBus after all modules
        if self.capture_bus:
            try:
                self.capture_bus.stop()
                self.capture_bus = None
            except Exception as e:
                logger.error("Failed to stop CaptureBus: %s", e)

    def get_status(self, name: str = None) -> dict:
        """Get status of one module or all modules."""
        if name:
            entry = self._modules.get(name)
            if entry is None:
                return {"error": "not registered"}
            alive = entry.process.is_alive() if entry.process else False
            return {
                "name": name,
                "running": alive,
                "pid": entry.process.pid if entry.process else None,
                "uptime": time.time() - entry.started_at if entry.started_at and alive else 0,
                "type": entry.module_class.module_type,
                "priority": entry.module_class.priority,
            }

        result = {}
        for n in self._modules:
            result[n] = self.get_status(n)
        return result

    # ------------------------------------------------------------------
    # Gating / enforcement
    # ------------------------------------------------------------------

    def _phase_allows(self, cls: Type[BaseModule]) -> bool:
        """Check if current engagement phase allows this module type."""
        if self._phase == "passive_only":
            return cls.module_type in ("passive", "stealth", "intel", "connectivity")
        return True  # active_allowed or unrestricted

    def _tier_allows(self, cls: Type[BaseModule]) -> bool:
        """Check if hardware tier limits allow this module."""
        limits = self._tier_limits
        running_by_type = self._count_running_by_type()

        if cls.module_type == "passive":
            if running_by_type.get("passive", 0) >= limits["max_passive"]:
                return False
        elif cls.module_type == "active":
            if running_by_type.get("active", 0) >= limits["max_active"]:
                return False
        elif cls.module_type == "stealth":
            if running_by_type.get("stealth", 0) >= limits.get("max_stealth", 99):
                return False
        elif cls.module_type in ("intel", "connectivity"):
            other_running = running_by_type.get("intel", 0) + running_by_type.get("connectivity", 0)
            if other_running >= limits.get("max_other", 99):
                return False
        return True

    def _count_running_by_type(self) -> dict[str, int]:
        """Count running modules by type."""
        counts: dict[str, int] = defaultdict(int)
        for entry in self._modules.values():
            if entry.process and entry.process.is_alive():
                counts[entry.module_class.module_type] += 1
        return dict(counts)

    def set_phase(self, phase: str) -> None:
        """Switch engagement phase (passive_only | active_allowed)."""
        self._phase = phase
        logger.info("Engagement phase set to: %s", phase)

    def check_scope(self, target: str) -> bool:
        """Check if a target IP/hostname is within engagement scope."""
        if not self._scope.get("enabled", False):
            return True

        import ipaddress

        networks = self._scope.get("networks", [])
        hosts = self._scope.get("hosts", [])
        exclude = self._scope.get("exclude", [])

        # Check exclusions first
        try:
            addr = ipaddress.ip_address(target)
            for exc in exclude:
                if addr in ipaddress.ip_network(exc, strict=False):
                    return False
            for net in networks:
                if addr in ipaddress.ip_network(net, strict=False):
                    return True
        except ValueError:
            pass  # Not an IP, check hostname

        if target in hosts:
            return True

        domains = self._scope.get("domains", [])
        for dom in domains:
            if target.endswith("." + dom) or target == dom:
                return True

        return False

    @property
    def tier(self) -> str:
        return self._tier

    @property
    def tier_limits(self) -> dict:
        return dict(self._tier_limits)

    @property
    def phase(self) -> str:
        return self._phase
