#!/usr/bin/env python3
"""VESC CAN health monitor node for SaltyBot.

Monitors both VESC motor controllers (CAN IDs 61 and 79) for frame drops,
bus-off errors, and timeouts.  Drives a recovery sequence on fault and
escalates to an e-stop if a VESC stays unresponsive beyond 2 s.

Subscribed topics
-----------------
  /vesc/left/state   (std_msgs/String) — JSON telemetry from vesc_telemetry node
  /vesc/right/state  (std_msgs/String) — JSON telemetry from vesc_telemetry node

Published topics
----------------
  /vesc/health   (std_msgs/String) — JSON health summary (10 Hz)
  /diagnostics   (diagnostic_msgs/DiagnosticArray)
  /estop         (std_msgs/String) — JSON e-stop event on state change
  /cmd_vel       (geometry_msgs/Twist) — zero velocity when e-stop active

Parameters
----------
  can_interface          str    'can0'  SocketCAN interface
  left_can_id            int    61
  right_can_id           int    79
  sender_can_id          int    127     This node's CAN ID for alive frames
  frame_timeout_s        float  0.5     Timeout before recovery starts
  estop_timeout_s        float  2.0     Timeout before e-stop escalation
  health_rate_hz         int    10      Publish + watchdog rate
  bus_off_check_rate_hz  int    1       How often to poll CAN interface state
"""

from __future__ import annotations

import json
import subprocess
import threading
import time
from typing import Optional

try:
    import rclpy
    from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue
    from geometry_msgs.msg import Twist
    from rclpy.node import Node
    from std_msgs.msg import String
    _ROS_AVAILABLE = True
except ImportError:  # pragma: no cover — missing outside ROS2 env
    _ROS_AVAILABLE = False
    Node = object  # type: ignore[assignment,misc]

try:
    import can as python_can
    _CAN_AVAILABLE = True
except ImportError:
    _CAN_AVAILABLE = False

from .recovery_fsm import Action, HealthFsm, VescHealthState


class VescHealthMonitor(Node):
    """ROS2 node: VESC CAN health watchdog + auto-recovery."""

    def __init__(self) -> None:
        super().__init__("vesc_health_monitor")

        # ----------------------------------------------------------------
        # Parameters
        # ----------------------------------------------------------------
        self.declare_parameter("can_interface",         "can0")
        self.declare_parameter("left_can_id",           61)
        self.declare_parameter("right_can_id",          79)
        self.declare_parameter("sender_can_id",         127)
        self.declare_parameter("frame_timeout_s",       0.5)
        self.declare_parameter("estop_timeout_s",       2.0)
        self.declare_parameter("health_rate_hz",        10)
        self.declare_parameter("bus_off_check_rate_hz", 1)

        self._iface     = self.get_parameter("can_interface").value
        self._left_id   = self.get_parameter("left_can_id").value
        self._right_id  = self.get_parameter("right_can_id").value
        self._sender_id = self.get_parameter("sender_can_id").value
        timeout_s       = self.get_parameter("frame_timeout_s").value
        estop_s         = self.get_parameter("estop_timeout_s").value
        hz              = max(1, int(self.get_parameter("health_rate_hz").value))
        bus_hz          = max(1, int(self.get_parameter("bus_off_check_rate_hz").value))

        # ----------------------------------------------------------------
        # State machine
        # ----------------------------------------------------------------
        self._fsm = HealthFsm(
            left_id=self._left_id,
            right_id=self._right_id,
            timeout_s=timeout_s,
            estop_s=estop_s,
        )
        self._fsm_lock = threading.Lock()

        # ----------------------------------------------------------------
        # Fault event log (ring buffer, newest last, max 200 entries)
        # ----------------------------------------------------------------
        self._fault_log: list[dict] = []
        self._fault_log_lock = threading.Lock()

        # ----------------------------------------------------------------
        # SocketCAN (for sending alive / GET_VALUES recovery frames)
        # ----------------------------------------------------------------
        self._bus: Optional["python_can.BusABC"] = None
        self._init_can()

        # ----------------------------------------------------------------
        # Publishers
        # ----------------------------------------------------------------
        self._pub_health = self.create_publisher(String,         "/vesc/health",  10)
        self._pub_diag   = self.create_publisher(DiagnosticArray, "/diagnostics",  10)
        self._pub_estop  = self.create_publisher(String,         "/estop",        10)
        self._pub_cmd    = self.create_publisher(Twist,          "/cmd_vel",      10)

        # ----------------------------------------------------------------
        # Subscribers
        # ----------------------------------------------------------------
        self.create_subscription(String, "/vesc/left/state",  self._on_left_state,  10)
        self.create_subscription(String, "/vesc/right/state", self._on_right_state, 10)

        # ----------------------------------------------------------------
        # Timers
        # ----------------------------------------------------------------
        self.create_timer(1.0 / hz,     self._on_health_tick)
        self.create_timer(1.0 / bus_hz, self._check_bus_off)

        self.get_logger().info(
            f"vesc_health_monitor: iface={self._iface}, "
            f"left={self._left_id}, right={self._right_id}, "
            f"timeout={timeout_s}s, estop={estop_s}s, {hz} Hz"
        )

    # ----------------------------------------------------------------
    # CAN initialisation
    # ----------------------------------------------------------------

    def _init_can(self) -> None:
        if not _CAN_AVAILABLE:
            self.get_logger().warn(
                "python-can not installed — alive frame recovery disabled"
            )
            return
        try:
            self._bus = python_can.interface.Bus(
                channel=self._iface, bustype="socketcan"
            )
            self.get_logger().info(f"CAN bus opened for recovery: {self._iface}")
        except Exception as exc:
            self.get_logger().warn(
                f"Could not open CAN bus for recovery frames ({exc}); "
                "health monitoring continues via topic subscription"
            )

    # ----------------------------------------------------------------
    # State subscribers
    # ----------------------------------------------------------------

    def _on_left_state(self, msg: String) -> None:
        self._handle_state_msg(msg, self._left_id)

    def _on_right_state(self, msg: String) -> None:
        self._handle_state_msg(msg, self._right_id)

    def _handle_state_msg(self, msg: String, can_id: int) -> None:
        try:
            d = json.loads(msg.data)
        except json.JSONDecodeError:
            return

        # The telemetry node sets "alive": true when the VESC is responding.
        # We track our own timeout independently for tighter control.
        now = time.monotonic()
        if d.get("alive", False):
            with self._fsm_lock:
                actions = self._fsm.on_frame(can_id, now)
            self._dispatch_actions(actions, can_id)

    # ----------------------------------------------------------------
    # Health timer (main watchdog loop)
    # ----------------------------------------------------------------

    def _on_health_tick(self) -> None:
        now = time.monotonic()
        with self._fsm_lock:
            per_vesc = self._fsm.tick(now)

        for can_id, actions in per_vesc.items():
            self._dispatch_actions(actions, can_id)

        self._publish_health(now)
        self._publish_diagnostics(now)

        # If any VESC is in e-stop, keep asserting zero velocity
        with self._fsm_lock:
            if self._fsm.any_estop:
                self._pub_cmd.publish(Twist())  # zero Twist = stop

    # ----------------------------------------------------------------
    # Bus-off check
    # ----------------------------------------------------------------

    def _check_bus_off(self) -> None:
        bus_off = self._is_can_bus_off()
        now = time.monotonic()
        with self._fsm_lock:
            if bus_off:
                actions = self._fsm.on_bus_off(now)
            else:
                actions = self._fsm.on_bus_ok(now)
        self._dispatch_actions(actions, can_id=None)

    def _is_can_bus_off(self) -> bool:
        """Return True if the CAN interface is in bus-off state."""
        try:
            result = subprocess.run(
                ["ip", "-json", "link", "show", self._iface],
                capture_output=True, text=True, timeout=2.0
            )
            if result.returncode != 0:
                return False
            data = json.loads(result.stdout)
            if not data:
                return False
            flags = data[0].get("flags", [])
            # Linux SocketCAN sets "BUS-OFF" in the flags list
            return "BUS-OFF" in [f.upper() for f in flags]
        except Exception:
            return False

    # ----------------------------------------------------------------
    # Action dispatcher
    # ----------------------------------------------------------------

    def _dispatch_actions(self, actions: list[Action], can_id: int | None) -> None:
        label = self._id_to_label(can_id) if can_id is not None else "bus"
        for action in actions:
            if action == Action.SEND_ALIVE:
                self._send_alive(can_id)
            elif action == Action.TRIGGER_ESTOP:
                self._on_estop_trigger(can_id, label)
            elif action == Action.CLEAR_ESTOP:
                self._on_estop_clear(can_id, label)
            elif action == Action.LOG_WARN:
                self.get_logger().warn(
                    f"VESC {label} (id={can_id}): "
                    "no CAN frames for >500 ms — sending alive frame"
                )
                self._log_fault_event("timeout", can_id, label)
            elif action == Action.LOG_ERROR:
                self.get_logger().error(
                    f"VESC {label} (id={can_id}): "
                    "unresponsive >2 s — e-stop triggered"
                )
                self._log_fault_event("estop", can_id, label)
            elif action == Action.LOG_RECOVERY:
                self.get_logger().info(
                    f"VESC {label} (id={can_id}): CAN frames resumed — recovered"
                )
                self._log_fault_event("recovery", can_id, label)

    def _on_estop_trigger(self, can_id: int | None, label: str) -> None:
        self._pub_cmd.publish(Twist())  # immediate zero velocity
        payload = json.dumps({
            "event": "estop_triggered",
            "can_id": can_id,
            "label": label,
            "stamp": time.time(),
        })
        self._pub_estop.publish(String(data=payload))

    def _on_estop_clear(self, can_id: int | None, label: str) -> None:
        payload = json.dumps({
            "event": "estop_cleared",
            "can_id": can_id,
            "label": label,
            "stamp": time.time(),
        })
        self._pub_estop.publish(String(data=payload))

    # ----------------------------------------------------------------
    # Alive / recovery frame
    # ----------------------------------------------------------------

    def _send_alive(self, can_id: int | None) -> None:
        """Send GET_VALUES request to prompt telemetry refresh."""
        if self._bus is None or can_id is None:
            return
        try:
            arb_id, payload = _make_get_values_request(self._sender_id, can_id)
            self._bus.send(python_can.Message(
                arbitration_id=arb_id,
                data=payload,
                is_extended_id=True,
            ))
            self.get_logger().debug(
                f"Sent GET_VALUES alive frame to VESC id={can_id}"
            )
        except Exception as exc:
            self.get_logger().warn(f"Failed to send alive frame (id={can_id}): {exc}")

    # ----------------------------------------------------------------
    # Publishers
    # ----------------------------------------------------------------

    def _publish_health(self, now: float) -> None:
        with self._fsm_lock:
            left_state   = self._fsm.left.state.value
            right_state  = self._fsm.right.state.value
            left_elapsed  = self._fsm.left.elapsed(now)
            right_elapsed = self._fsm.right.elapsed(now)
            left_estop    = self._fsm.left.estop_active
            right_estop   = self._fsm.right.estop_active
            bus_off       = self._fsm._bus_off
            any_estop     = self._fsm.any_estop

        health = {
            "stamp": time.time(),
            "left": {
                "can_id":       self._left_id,
                "state":        left_state,
                "elapsed_s":    round(min(left_elapsed, 9999.0), 3),
                "estop_active": left_estop,
            },
            "right": {
                "can_id":       self._right_id,
                "state":        right_state,
                "elapsed_s":    round(min(right_elapsed, 9999.0), 3),
                "estop_active": right_estop,
            },
            "bus_off":     bus_off,
            "estop_active": any_estop,
        }

        with self._fault_log_lock:
            health["recent_faults"] = list(self._fault_log[-10:])

        self._pub_health.publish(String(data=json.dumps(health)))

    def _publish_diagnostics(self, now: float) -> None:
        array = DiagnosticArray()
        array.header.stamp = self.get_clock().now().to_msg()

        with self._fsm_lock:
            monitors = [
                (self._fsm.left,  "left"),
                (self._fsm.right, "right"),
            ]

        for mon, label in monitors:
            status = DiagnosticStatus()
            status.name        = f"VESC/health/{label} (CAN ID {mon.can_id})"
            status.hardware_id = f"vesc_health_{mon.can_id}"

            elapsed = mon.elapsed(now)

            if mon.state == VescHealthState.ESTOP:
                status.level   = DiagnosticStatus.ERROR
                status.message = "E-STOP: VESC unresponsive >2 s"
            elif mon.state == VescHealthState.BUS_OFF:
                status.level   = DiagnosticStatus.ERROR
                status.message = "CAN bus-off error"
            elif mon.state == VescHealthState.DEGRADED:
                status.level   = DiagnosticStatus.WARN
                status.message = f"Frame timeout ({elapsed:.2f} s) — recovery active"
            else:
                status.level   = DiagnosticStatus.OK
                status.message = "OK"

            status.values = [
                KeyValue(key="state",       value=mon.state.value),
                KeyValue(key="elapsed_s",   value=f"{min(elapsed, 9999.0):.3f}"),
                KeyValue(key="estop_active", value=str(mon.estop_active)),
            ]
            array.status.append(status)

        self._pub_diag.publish(array)

    # ----------------------------------------------------------------
    # Fault event log
    # ----------------------------------------------------------------

    def _log_fault_event(self, event: str, can_id: int | None, label: str) -> None:
        entry = {
            "event":  event,
            "can_id": can_id,
            "label":  label,
            "stamp":  time.time(),
        }
        with self._fault_log_lock:
            self._fault_log.append(entry)
            if len(self._fault_log) > 200:
                self._fault_log = self._fault_log[-200:]

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

    def _id_to_label(self, can_id: int | None) -> str:
        if can_id == self._left_id:
            return "left"
        if can_id == self._right_id:
            return "right"
        return str(can_id)

    # ----------------------------------------------------------------
    # Cleanup
    # ----------------------------------------------------------------

    def destroy_node(self) -> None:
        if self._bus is not None:
            self._bus.shutdown()
        super().destroy_node()


# ---------------------------------------------------------------------------
# Inline GET_VALUES request builder (avoids cross-package import at runtime)
# ---------------------------------------------------------------------------

_CAN_PACKET_PROCESS_SHORT_BUFFER = 32
_COMM_GET_VALUES = 4


def _make_get_values_request(sender_id: int, target_id: int) -> tuple[int, bytes]:
    """Build a GET_VALUES short-buffer CAN frame for the given VESC."""
    arb_id  = (_CAN_PACKET_PROCESS_SHORT_BUFFER << 8) | (target_id & 0xFF)
    payload = bytes([sender_id & 0xFF, 0x00, _COMM_GET_VALUES])
    return arb_id, payload


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

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


if __name__ == "__main__":
    main()
