#!/usr/bin/env python3
"""Watchdog module — monitor all SystemMonitor modules and managed tools.

Performs periodic health checks, auto-restarts crashed modules (max 3 attempts),
and manages OOM priority assignments.  Publishes MODULE_RESTARTED events on
successful recovery.
"""

import logging
import os
import signal
import threading
import time
from typing import Dict, Optional

from modules.base import BaseModule
from core.bus import Event
from utils.resource import get_process_resources

logger = logging.getLogger(__name__)

# OOM priority assignments by module type
_OOM_PRIORITIES = {
    "core": -500,
    "stealth": -400,
    "connectivity": -300,
    "intel": -200,
    "passive": 100,
    "active": 200,
}

MAX_RESTART_ATTEMPTS = 3
HEALTH_CHECK_INTERVAL = 30.0  # seconds


class _ModuleEntry:
    """Internal tracking for a monitored module."""
    __slots__ = ("name", "module_type", "pid", "restart_count", "last_check",
                 "last_status", "failed")

    def __init__(self, name: str, module_type: str, pid: int):
        self.name = name
        self.module_type = module_type
        self.pid = pid
        self.restart_count = 0
        self.last_check = 0.0
        self.last_status = "unknown"
        self.failed = False


class Watchdog(BaseModule):
    """Monitor modules and tools, auto-restart on failure, manage OOM priorities."""

    name = "watchdog"
    module_type = "stealth"
    priority = -500
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._monitored: Dict[str, _ModuleEntry] = {}
        self._lock = threading.Lock()
        self._check_thread: Optional[threading.Thread] = None
        self._check_interval = HEALTH_CHECK_INTERVAL

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        # Subscribe to error/crash events
        self.bus.subscribe(self._on_module_error, event_type="MODULE_ERROR")
        self.bus.subscribe(self._on_tool_crashed, event_type="TOOL_CRASHED")
        self.bus.subscribe(self._on_module_started, event_type="MODULE_STARTED")
        self.bus.subscribe(self._on_module_stopped, event_type="MODULE_STOPPED")

        # Populate initial monitoring list from state manager
        self._refresh_monitored_list()

        # Start health check loop
        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self._check_thread = threading.Thread(
            target=self._health_check_loop, daemon=True, name="sensor-watchdog",
        )
        self._check_thread.start()

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("Watchdog active — monitoring %d modules", len(self._monitored))

    def stop(self) -> None:
        if not self._running:
            return

        self._running = False

        self.bus.unsubscribe(self._on_module_error, event_type="MODULE_ERROR")
        self.bus.unsubscribe(self._on_tool_crashed, event_type="TOOL_CRASHED")
        self.bus.unsubscribe(self._on_module_started, event_type="MODULE_STARTED")
        self.bus.unsubscribe(self._on_module_stopped, event_type="MODULE_STOPPED")

        if self._check_thread and self._check_thread.is_alive():
            self._check_thread.join(timeout=5.0)

        self.state.set_module_status(self.name, "stopped")
        logger.info("Watchdog stopped")

    def status(self) -> dict:
        with self._lock:
            health_report = []
            for name, entry in self._monitored.items():
                health_report.append({
                    "module": entry.name,
                    "module_type": entry.module_type,
                    "pid": entry.pid,
                    "status": entry.last_status,
                    "restarts": entry.restart_count,
                    "last_check": entry.last_check,
                    "failed": entry.failed,
                })

        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "monitored_count": len(self._monitored),
            "health_report": health_report,
            "check_interval_s": self._check_interval,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        self._check_interval = config.get("check_interval_s", HEALTH_CHECK_INTERVAL)

    # ------------------------------------------------------------------
    # Event handlers
    # ------------------------------------------------------------------

    def _on_module_started(self, event: Event) -> None:
        """Register a newly started module for monitoring."""
        payload = event.payload
        name = payload.get("module", "")
        pid = payload.get("pid")
        module_type = payload.get("module_type", "passive")

        if name and pid and name != self.name:
            with self._lock:
                self._monitored[name] = _ModuleEntry(name, module_type, pid)
            self._set_oom_priority(pid, module_type)
            logger.debug("Now monitoring: %s (PID %d, type %s)", name, pid, module_type)

    def _on_module_stopped(self, event: Event) -> None:
        """Remove a stopped module from monitoring."""
        name = event.payload.get("module", "")
        if name:
            with self._lock:
                self._monitored.pop(name, None)

    def _on_module_error(self, event: Event) -> None:
        """Handle a module error event — attempt restart."""
        name = event.payload.get("module", "")
        error = event.payload.get("error", "unknown")
        logger.warning("Module error reported: %s — %s", name, error)
        self._attempt_restart(name)

    def _on_tool_crashed(self, event: Event) -> None:
        """Handle a managed tool crash — delegate restart to engine/tool_manager."""
        tool = event.payload.get("tool", "")
        pid = event.payload.get("pid")
        logger.warning("Tool crashed: %s (PID %s)", tool, pid)

        # Tool restarts are handled by tool_manager; we just log and track
        with self._lock:
            entry = self._monitored.get(tool)
            if entry:
                entry.restart_count += 1
                entry.last_status = "crashed"

    # ------------------------------------------------------------------
    # Health check loop
    # ------------------------------------------------------------------

    def _health_check_loop(self) -> None:
        """Periodic health check of all monitored modules."""
        while self._running:
            time.sleep(self._check_interval)
            if not self._running:
                break

            self._run_health_checks()

    def _run_health_checks(self) -> None:
        """Check health of all monitored modules."""
        with self._lock:
            entries = list(self._monitored.values())

        for entry in entries:
            if entry.failed:
                continue  # Already gave up on this module

            now = time.time()
            entry.last_check = now

            # Check if process is alive
            if not self._is_pid_alive(entry.pid):
                logger.warning("Module %s (PID %d) is dead", entry.name, entry.pid)
                entry.last_status = "dead"
                self._attempt_restart(entry.name)
                continue

            # Check process health via engine if available
            if self.engine:
                try:
                    module_instance = self._get_module_instance(entry.name)
                    if module_instance and hasattr(module_instance, "health_check"):
                        healthy = module_instance.health_check()
                        entry.last_status = "healthy" if healthy else "unhealthy"
                        if not healthy:
                            logger.warning("Module %s failed health check", entry.name)
                            self._attempt_restart(entry.name)
                            continue
                    else:
                        entry.last_status = "alive"
                except Exception as exc:
                    logger.debug("Health check error for %s: %s", entry.name, exc)
                    entry.last_status = "check_error"
            else:
                # No engine reference — just check PID is alive
                entry.last_status = "alive"

    def _attempt_restart(self, module_name: str) -> None:
        """Attempt to restart a failed module via the engine."""
        with self._lock:
            entry = self._monitored.get(module_name)
            if not entry:
                return
            if entry.failed:
                return

            entry.restart_count += 1
            if entry.restart_count > MAX_RESTART_ATTEMPTS:
                entry.failed = True
                entry.last_status = "given_up"
                logger.error(
                    "Module %s exceeded max restarts (%d) — giving up",
                    module_name, MAX_RESTART_ATTEMPTS,
                )
                self.bus.emit(
                    "MODULE_ERROR",
                    {"module": module_name, "error": "max_restarts_exceeded",
                     "restarts": entry.restart_count},
                    source_module=self.name,
                )
                return

        logger.info(
            "Attempting restart of %s (attempt %d/%d)",
            module_name, entry.restart_count, MAX_RESTART_ATTEMPTS,
        )

        # Delegate restart to engine
        if self.engine and hasattr(self.engine, "restart_module"):
            try:
                success = self.engine.restart_module(module_name)
                if success:
                    entry.last_status = "restarted"
                    self.bus.emit(
                        "MODULE_RESTARTED" if "MODULE_RESTARTED" in getattr(
                            self.bus, '_subscribers', {}
                        ) else "TOOL_RESTARTED",
                        {"module": module_name, "attempt": entry.restart_count},
                        source_module=self.name,
                    )
                    # Publish MODULE_RESTARTED (it may not be in EVENT_TYPES but
                    # subscribers can still listen for it)
                    self.bus.emit(
                        "TOOL_RESTARTED",
                        {"module": module_name, "tool": module_name,
                         "attempt": entry.restart_count},
                        source_module=self.name,
                    )
                    logger.info("Successfully restarted %s", module_name)
                else:
                    entry.last_status = "restart_failed"
                    logger.error("Engine failed to restart %s", module_name)
            except Exception as exc:
                entry.last_status = "restart_error"
                logger.error("Restart error for %s: %s", module_name, exc)
        else:
            logger.warning("No engine available — cannot restart %s", module_name)
            entry.last_status = "no_engine"

    # ------------------------------------------------------------------
    # OOM priority management
    # ------------------------------------------------------------------

    def _set_oom_priority(self, pid: int, module_type: str) -> None:
        """Set OOM score adjustment for a process based on its module type."""
        score = _OOM_PRIORITIES.get(module_type, 0)
        oom_path = f"/proc/{pid}/oom_score_adj"
        try:
            with open(oom_path, "w") as f:
                f.write(str(score))
            logger.debug("Set OOM priority for PID %d: %d (%s)", pid, score, module_type)
        except (IOError, PermissionError) as exc:
            logger.debug("Cannot set OOM priority for PID %d: %s", pid, exc)

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

    def _refresh_monitored_list(self) -> None:
        """Populate monitored list from state manager's module status."""
        try:
            all_status = self.state.get_all_module_status()
            with self._lock:
                for name, info in all_status.items():
                    if name == self.name:
                        continue
                    if info.get("status") == "running" and info.get("pid"):
                        self._monitored[name] = _ModuleEntry(
                            name=name,
                            module_type="unknown",
                            pid=info["pid"],
                        )
        except Exception as exc:
            logger.debug("Could not refresh monitored list: %s", exc)

    def _get_module_instance(self, module_name: str):
        """Get a module instance from the engine if available."""
        if self.engine and hasattr(self.engine, "get_module"):
            return self.engine.get_module(module_name)
        return None

    @staticmethod
    def _is_pid_alive(pid: int) -> bool:
        """Check if a process with the given PID exists."""
        try:
            os.kill(pid, 0)
            return True
        except (ProcessLookupError, PermissionError):
            return False
        except OSError:
            return False
