#!/usr/bin/env python3
"""Node watchdog monitor for SaltyBot critical systems.

Monitors heartbeats from balance, motor driver, emergency, and docking nodes.
Publishes alerts on heartbeat loss >1s. Implements safety fallback: zeros cmd_vel
if motor driver lost >2s.

Published topics:
  /saltybot/node_watchdog (std_msgs/String) - JSON watchdog status
  /saltybot/cmd_vel_safe (geometry_msgs/Twist) - cmd_vel with motor driver safety check

Subscribed topics:
  /saltybot/balance_heartbeat (std_msgs/UInt32) - Balance node heartbeat
  /saltybot/motor_driver_heartbeat (std_msgs/UInt32) - Motor driver heartbeat
  /saltybot/emergency_heartbeat (std_msgs/UInt32) - Emergency system heartbeat
  /saltybot/docking_heartbeat (std_msgs/UInt32) - Docking node heartbeat
  /cmd_vel (geometry_msgs/Twist) - Velocity command input
"""

import json
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from geometry_msgs.msg import Twist
from std_msgs.msg import UInt32, String


class NodeWatchdogNode(Node):
    """ROS2 watchdog monitor for critical system nodes."""

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

        # Parameters
        self.declare_parameter("frequency", 20)  # Hz
        self.declare_parameter("heartbeat_timeout", 1.0)  # seconds, general timeout
        self.declare_parameter("motor_driver_critical_timeout", 2.0)  # seconds

        frequency = self.get_parameter("frequency").value
        self.heartbeat_timeout = self.get_parameter("heartbeat_timeout").value
        self.motor_driver_critical_timeout = self.get_parameter(
            "motor_driver_critical_timeout"
        ).value

        # Heartbeat tracking
        self.critical_nodes = {
            "balance": None,
            "motor_driver": None,
            "emergency": None,
            "docking": None,
        }
        self.last_heartbeat_time = {
            "balance": None,
            "motor_driver": None,
            "emergency": None,
            "docking": None,
        }
        self.last_cmd_vel = None
        self.motor_driver_down = False

        # Subscriptions for heartbeats
        self.create_subscription(
            UInt32, "/saltybot/balance_heartbeat", self._on_balance_heartbeat, 10
        )
        self.create_subscription(
            UInt32,
            "/saltybot/motor_driver_heartbeat",
            self._on_motor_driver_heartbeat,
            10,
        )
        self.create_subscription(
            UInt32, "/saltybot/emergency_heartbeat", self._on_emergency_heartbeat, 10
        )
        self.create_subscription(
            UInt32, "/saltybot/docking_heartbeat", self._on_docking_heartbeat, 10
        )

        # cmd_vel subscription and safe republishing
        self.create_subscription(Twist, "/cmd_vel", self._on_cmd_vel, 10)

        # Publications
        self.pub_watchdog = self.create_publisher(String, "/saltybot/node_watchdog", 10)
        self.pub_cmd_vel_safe = self.create_publisher(
            Twist, "/saltybot/cmd_vel_safe", 10
        )

        # Timer for periodic monitoring at 20Hz
        period = 1.0 / frequency
        self.timer: Timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Node watchdog initialized at {frequency}Hz. "
            f"Heartbeat timeout: {self.heartbeat_timeout}s, "
            f"Motor driver critical: {self.motor_driver_critical_timeout}s"
        )

    def _on_balance_heartbeat(self, msg: UInt32) -> None:
        """Update balance node heartbeat timestamp."""
        self.last_heartbeat_time["balance"] = self.get_clock().now()

    def _on_motor_driver_heartbeat(self, msg: UInt32) -> None:
        """Update motor driver heartbeat timestamp."""
        self.last_heartbeat_time["motor_driver"] = self.get_clock().now()
        self.motor_driver_down = False

    def _on_emergency_heartbeat(self, msg: UInt32) -> None:
        """Update emergency system heartbeat timestamp."""
        self.last_heartbeat_time["emergency"] = self.get_clock().now()

    def _on_docking_heartbeat(self, msg: UInt32) -> None:
        """Update docking node heartbeat timestamp."""
        self.last_heartbeat_time["docking"] = self.get_clock().now()

    def _on_cmd_vel(self, msg: Twist) -> None:
        """Cache the last received cmd_vel."""
        self.last_cmd_vel = msg

    def _check_node_health(self) -> dict:
        """Check health of all monitored nodes.

        Returns:
            dict: Health status of each node with timeout and elapsed time.
        """
        now = self.get_clock().now()
        health = {}

        for node_name in self.critical_nodes:
            last_time = self.last_heartbeat_time[node_name]

            if last_time is None:
                # No heartbeat received yet
                health[node_name] = {
                    "status": "unknown",
                    "elapsed_s": None,
                    "timeout_s": self.heartbeat_timeout,
                }
            else:
                # Calculate elapsed time since last heartbeat
                elapsed = (now - last_time).nanoseconds / 1e9
                is_timeout = elapsed > self.heartbeat_timeout

                # Special case: motor driver has longer critical timeout
                if node_name == "motor_driver":
                    is_critical = elapsed > self.motor_driver_critical_timeout
                else:
                    is_critical = False

                health[node_name] = {
                    "status": "down" if is_timeout else "up",
                    "elapsed_s": elapsed,
                    "timeout_s": (
                        self.motor_driver_critical_timeout
                        if node_name == "motor_driver"
                        else self.heartbeat_timeout
                    ),
                    "critical": is_critical,
                }

        return health

    def _timer_callback(self) -> None:
        """Monitor node health and publish alerts at 20Hz."""
        health = self._check_node_health()

        # Detect if motor driver is in critical state (down for >2s)
        motor_driver_health = health.get("motor_driver", {})
        if motor_driver_health.get("critical", False):
            self.motor_driver_down = True
            self.get_logger().warn(
                f"MOTOR DRIVER DOWN >2s ({motor_driver_health['elapsed_s']:.1f}s). "
                "Applying safety fallback: zeroing cmd_vel."
            )

        # Determine any nodes down for >1s
        nodes_with_timeout = {
            name: status
            for name, status in health.items()
            if status.get("status") == "down"
        }

        # Publish watchdog status
        watchdog_status = {
            "timestamp": self.get_clock().now().nanoseconds / 1e9,
            "all_healthy": len(nodes_with_timeout) == 0
            and not self.motor_driver_down,
            "health": health,
            "motor_driver_critical": self.motor_driver_down,
        }
        watchdog_msg = String(data=json.dumps(watchdog_status))
        self.pub_watchdog.publish(watchdog_msg)

        # Publish cmd_vel with safety checks
        if self.last_cmd_vel is not None:
            cmd_vel_safe = self._apply_safety_checks(self.last_cmd_vel)
            self.pub_cmd_vel_safe.publish(cmd_vel_safe)

    def _apply_safety_checks(self, cmd_vel: Twist) -> Twist:
        """Apply safety checks to cmd_vel based on system state.

        Args:
            cmd_vel: Original velocity command

        Returns:
            Twist: Potentially modified velocity command for safe operation.
        """
        safe_cmd = Twist()

        # If motor driver is critically down, zero all velocities
        if self.motor_driver_down:
            return safe_cmd

        # Otherwise, pass through unchanged
        safe_cmd.linear.x = cmd_vel.linear.x
        safe_cmd.linear.y = cmd_vel.linear.y
        safe_cmd.linear.z = cmd_vel.linear.z
        safe_cmd.angular.x = cmd_vel.angular.x
        safe_cmd.angular.y = cmd_vel.angular.y
        safe_cmd.angular.z = cmd_vel.angular.z
        return safe_cmd


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


if __name__ == "__main__":
    main()
