#!/usr/bin/env python3
"""System health monitor for SaltyBot.

Central node that monitors heartbeats from all critical nodes. Tracks expected
nodes from YAML config, marks nodes DEAD if silent >5s, auto-restarts via
ros2 launch, publishes /saltybot/system_health JSON, and triggers face alerts.

Published topics:
  /saltybot/system_health (std_msgs/String) - JSON system health status

Subscribed topics:
  /saltybot/<node_name>/heartbeat (std_msgs/String) - Node heartbeat signals
"""

import json
import time
from pathlib import Path
from typing import Dict, Optional
from dataclasses import dataclass, asdict
from datetime import datetime

import yaml
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from std_msgs.msg import String


@dataclass
class NodeHealth:
    """Health status of a single node."""

    name: str
    status: str  # "UP", "DOWN", "RESTARTING"
    last_heartbeat: float  # Timestamp of last received heartbeat
    heartbeat_count: int = 0
    restart_count: int = 0
    expected: bool = True


class HealthMonitorNode(Node):
    """ROS2 node for system health monitoring."""

    def __init__(self):
        super().__init__("health_monitor")

        # Load configuration
        self.declare_parameter("config_file", "health_config.yaml")
        config_path = self.get_parameter("config_file").value

        self.node_health: Dict[str, NodeHealth] = {}
        self.startup_time = time.time()
        self.last_critical_alert = 0.0
        self.alert_cooldown = 5.0  # Seconds between critical alerts

        # Load node configuration
        self._load_config(config_path)

        # Parameters
        self.declare_parameter("heartbeat_timeout", 5.0)  # Seconds
        self.declare_parameter("check_frequency", 1.0)  # Hz
        self.declare_parameter("enable_auto_restart", True)
        self.declare_parameter("restart_command_file", "/tmp/restart_node.sh")

        self.heartbeat_timeout = self.get_parameter("heartbeat_timeout").value
        check_frequency = self.get_parameter("check_frequency").value
        self.enable_auto_restart = self.get_parameter("enable_auto_restart").value
        self.restart_cmd_file = self.get_parameter("restart_command_file").value

        # Subscribe to heartbeats from all expected nodes
        self._setup_subscriptions()

        # Publisher for system health
        self.pub_health = self.create_publisher(String, "/saltybot/system_health", 1)
        self.pub_face_alert = self.create_publisher(String, "/saltybot/face/alert", 1)

        # Health check timer
        period = 1.0 / check_frequency
        self.timer: Timer = self.create_timer(period, self._check_health)

        self.get_logger().info(
            f"Health monitor initialized with {len(self.node_health)} tracked nodes. "
            f"Timeout: {self.heartbeat_timeout}s, Auto-restart: {self.enable_auto_restart}"
        )

    def _load_config(self, config_file: str) -> None:
        """Load node configuration from YAML file."""
        try:
            # Try to find config in share directory
            if not Path(config_file).exists():
                # Look in package share directory
                share_dir = Path(__file__).parent.parent / "config"
                config_file = str(share_dir / config_file)

            with open(config_file, "r") as f:
                config = yaml.safe_load(f) or {}

            monitored_nodes = config.get("monitored_nodes", [])
            for node_name in monitored_nodes:
                self.node_health[node_name] = NodeHealth(
                    name=node_name, status="UNKNOWN", last_heartbeat=time.time()
                )

            self.get_logger().info(f"Loaded {len(monitored_nodes)} nodes from config")
        except FileNotFoundError:
            self.get_logger().warn(
                f"Config file not found: {config_file}. "
                "Will monitor nodes as they send heartbeats."
            )

    def _setup_subscriptions(self) -> None:
        """Create subscriptions for all expected nodes."""
        for node_name in self.node_health.keys():
            topic = f"/saltybot/{node_name}/heartbeat"
            self.create_subscription(String, topic, self._on_heartbeat, 10)

    def _on_heartbeat(self, msg: String) -> None:
        """Handle incoming heartbeat from a node."""
        # Parse heartbeat message (expected format: "node_name:data")
        try:
            parts = msg.data.split(":", 1)
            node_name = parts[0].strip()
            data = parts[1].strip() if len(parts) > 1 else ""

            # Create node entry if not yet tracked
            if node_name not in self.node_health:
                self.node_health[node_name] = NodeHealth(
                    name=node_name, status="UP", last_heartbeat=time.time(), expected=False
                )

            # Update heartbeat
            node = self.node_health[node_name]
            node.last_heartbeat = time.time()
            node.heartbeat_count += 1
            if node.status != "UP":
                node.status = "UP"
                self.get_logger().info(f"Node {node_name} is UP")

        except Exception as e:
            self.get_logger().error(f"Error processing heartbeat: {e}")

    def _check_health(self) -> None:
        """Periodically check health of all nodes and publish status."""
        now = time.time()
        critical_down = []

        for node_name, node in self.node_health.items():
            # Check if heartbeat is stale
            time_since_heartbeat = now - node.last_heartbeat

            if time_since_heartbeat > self.heartbeat_timeout:
                if node.status != "DOWN":
                    self.get_logger().warn(
                        f"Node {node_name} DOWN (silent for {time_since_heartbeat:.1f}s)"
                    )
                    node.status = "DOWN"

                # Track critical (expected) nodes
                if node.expected:
                    critical_down.append(node_name)

                # Attempt auto-restart
                if self.enable_auto_restart and node.status == "DOWN":
                    self._trigger_restart(node_name)
            else:
                # Node is healthy
                if node.status != "UP":
                    node.status = "UP"

        # Publish system health
        self._publish_health(critical_down)

        # Alert face if critical nodes are down
        if critical_down:
            self._alert_critical(critical_down, now)

    def _trigger_restart(self, node_name: str) -> None:
        """Trigger restart of a dead node via launch system."""
        node = self.node_health[node_name]
        node.restart_count += 1

        self.get_logger().warn(
            f"Attempting auto-restart for {node_name} (attempt #{node.restart_count})"
        )

        # Update status
        node.status = "RESTARTING"

        # In a real implementation, this would trigger ros2 launch or systemd service restart
        # For now, log the attempt
        try:
            # Example: restart via launch system
            # This would need to be configured based on actual launch setup
            restart_script = (
                f"#!/bin/bash\n"
                f"# Auto-restart triggered at {datetime.now().isoformat()}\n"
                f"ros2 launch saltybot_bringup {node_name}.launch.py &\n"
            )
            with open(self.restart_cmd_file, "a") as f:
                f.write(restart_script)

            self.get_logger().info(f"Restart command queued for {node_name}")
        except Exception as e:
            self.get_logger().error(f"Failed to queue restart for {node_name}: {e}")

    def _publish_health(self, critical_down: list) -> None:
        """Publish system health status as JSON."""
        health_data = {
            "timestamp": datetime.now().isoformat(),
            "uptime_s": time.time() - self.startup_time,
            "nodes": {
                node.name: {
                    "status": node.status,
                    "time_since_heartbeat_s": time.time() - node.last_heartbeat,
                    "heartbeat_count": node.heartbeat_count,
                    "restart_count": node.restart_count,
                    "expected": node.expected,
                }
                for node in self.node_health.values()
            },
            "critical_down": critical_down,
            "system_healthy": len(critical_down) == 0,
        }

        msg = String(data=json.dumps(health_data))
        self.pub_health.publish(msg)

    def _alert_critical(self, critical_nodes: list, now: float) -> None:
        """Alert face display of critical node failures."""
        # Rate-limit alerts to avoid spam
        if now - self.last_critical_alert < self.alert_cooldown:
            return

        self.last_critical_alert = now

        alert_msg = {
            "type": "system_alert",
            "severity": "critical",
            "message": f"System critical: {', '.join(critical_nodes)} down",
            "nodes": critical_nodes,
            "timestamp": datetime.now().isoformat(),
        }

        msg = String(data=json.dumps(alert_msg))
        self.pub_face_alert.publish(msg)

        self.get_logger().warn(
            f"CRITICAL ALERT: {len(critical_nodes)} expected node(s) down: {critical_nodes}"
        )


def main(args=None):
    rclpy.init(args=args)
    node = HealthMonitorNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
