#!/usr/bin/env python3
"""VESC CAN telemetry node for SaltyBot dual FSESC 6.7 Pro (FW 6.6).

Listens on SocketCAN (can0) for VESC status broadcast frames from two VESCs
and publishes parsed telemetry to ROS2 topics. Also sends periodic GET_VALUES
requests to trigger fresh STATUS broadcasts.

Published topics
----------------
  /vesc/left/state   (std_msgs/String) — JSON telemetry for left VESC
  /vesc/right/state  (std_msgs/String) — JSON telemetry for right VESC
  /vesc/combined     (std_msgs/String) — voltage, total current, both RPMs
  /diagnostics       (diagnostic_msgs/DiagnosticArray)

Parameters
----------
  can_interface              str   'can0'   SocketCAN interface
  left_can_id                int   56       CAN ID of left VESC
  right_can_id               int   68       CAN ID of right VESC
  poll_rate_hz               int   20       Publish + request rate (10-50 Hz)
  sender_can_id              int   127      This node's CAN ID
  overcurrent_threshold_a    float 60.0     Fault alert threshold (A)
  overtemp_fet_threshold_c   float 80.0     FET overtemp alert threshold (°C)
  overtemp_motor_threshold_c float 100.0    Motor overtemp alert threshold (°C)
"""

import json
import threading
import time
from typing import Optional

import rclpy
from rclpy.node import Node
from std_msgs.msg import String
from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue

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

from .vesc_can_protocol import (
    CAN_PACKET_STATUS,
    CAN_PACKET_STATUS_4,
    CAN_PACKET_STATUS_5,
    VescState,
    make_get_values_request,
    parse_status,
    parse_status_4,
    parse_status_5,
)


class VescTelemetryNode(Node):
    """ROS2 node: VESC CAN telemetry via SocketCAN."""

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

        # ----------------------------------------------------------------
        # Parameters
        # ----------------------------------------------------------------
        self.declare_parameter("can_interface",              "can0")
        self.declare_parameter("left_can_id",                56)
        self.declare_parameter("right_can_id",               68)
        self.declare_parameter("poll_rate_hz",               20)
        self.declare_parameter("sender_can_id",              127)
        self.declare_parameter("overcurrent_threshold_a",    60.0)
        self.declare_parameter("overtemp_fet_threshold_c",   80.0)
        self.declare_parameter("overtemp_motor_threshold_c", 100.0)

        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
        hz               = int(max(10, min(50, self.get_parameter("poll_rate_hz").value)))
        self._sender_id  = self.get_parameter("sender_can_id").value
        self._oc_thresh  = self.get_parameter("overcurrent_threshold_a").value
        self._otf_thresh = self.get_parameter("overtemp_fet_threshold_c").value
        self._otm_thresh = self.get_parameter("overtemp_motor_threshold_c").value

        # ----------------------------------------------------------------
        # Per-VESC state (protected by _state_lock)
        # ----------------------------------------------------------------
        self._state: dict[int, VescState] = {
            self._left_id:  VescState(can_id=self._left_id),
            self._right_id: VescState(can_id=self._right_id),
        }
        self._state_lock = threading.Lock()

        # Track which faults we have already logged to avoid log spam
        self._fault_logged: dict[int, int] = {
            self._left_id: 0, self._right_id: 0
        }

        # ----------------------------------------------------------------
        # Publishers
        # ----------------------------------------------------------------
        self._pub_left     = self.create_publisher(String,        "/vesc/left/state",  10)
        self._pub_right    = self.create_publisher(String,        "/vesc/right/state", 10)
        self._pub_combined = self.create_publisher(String,        "/vesc/combined",    10)
        self._pub_diag     = self.create_publisher(DiagnosticArray, "/diagnostics",    10)

        # ----------------------------------------------------------------
        # SocketCAN
        # ----------------------------------------------------------------
        self._bus: Optional["can.BusABC"] = None
        self._rx_thread: Optional[threading.Thread] = None
        self._running = False
        self._init_can()

        # ----------------------------------------------------------------
        # Publish + poll timer
        # ----------------------------------------------------------------
        self.create_timer(1.0 / hz, self._on_timer)

        self.get_logger().info(
            f"vesc_telemetry: iface={self._iface}, "
            f"left_id={self._left_id}, right_id={self._right_id}, "
            f"poll={hz} Hz, sender_id={self._sender_id}"
        )

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

    def _init_can(self) -> None:
        if not _CAN_AVAILABLE:
            self.get_logger().error(
                "python-can not installed — run: pip install python-can"
            )
            return
        try:
            self._bus = can.interface.Bus(channel=self._iface, bustype="socketcan")
            self._running = True
            self._rx_thread = threading.Thread(
                target=self._rx_loop, name="vesc_can_rx", daemon=True
            )
            self._rx_thread.start()
            self.get_logger().info(f"SocketCAN opened: {self._iface}")
        except Exception as exc:
            self.get_logger().error(f"Failed to open {self._iface}: {exc}")

    # ----------------------------------------------------------------
    # Receive loop (background thread)
    # ----------------------------------------------------------------

    def _rx_loop(self) -> None:
        """Drain SocketCAN frames and update VESC state."""
        while self._running and self._bus is not None:
            try:
                msg = self._bus.recv(timeout=0.1)
            except Exception:
                continue
            if msg is None or not msg.is_extended_id:
                continue
            self._dispatch_frame(msg.arbitration_id, bytes(msg.data))

    def _dispatch_frame(self, arb_id: int, data: bytes) -> None:
        """Decode extended CAN ID and update the matching VescState."""
        packet_type = (arb_id >> 8) & 0xFFFF
        vesc_id     = arb_id & 0xFF

        if vesc_id not in self._state:
            return

        now = time.monotonic()
        with self._state_lock:
            s = self._state[vesc_id]

            if packet_type == CAN_PACKET_STATUS:
                parsed = parse_status(data)
                if parsed:
                    s.rpm             = parsed["rpm"]
                    s.current_a       = parsed["current_a"]
                    s.duty_cycle      = parsed["duty_cycle"]
                    s.last_status_ts  = now

            elif packet_type == CAN_PACKET_STATUS_4:
                parsed = parse_status_4(data)
                if parsed:
                    s.temp_fet_c   = parsed["temp_fet_c"]
                    s.temp_motor_c = parsed["temp_motor_c"]
                    s.current_in_a = parsed["current_in_a"]

            elif packet_type == CAN_PACKET_STATUS_5:
                parsed = parse_status_5(data)
                if parsed:
                    s.voltage_v        = parsed["voltage_v"]
                    s.last_status5_ts  = now

    # ----------------------------------------------------------------
    # Timer: poll + publish
    # ----------------------------------------------------------------

    def _on_timer(self) -> None:
        self._send_poll_requests()
        self._publish_all()

    def _send_poll_requests(self) -> None:
        """Send GET_VALUES to both VESCs to prompt STATUS broadcast refresh."""
        if self._bus is None:
            return
        for vesc_id in (self._left_id, self._right_id):
            try:
                arb_id, payload = make_get_values_request(self._sender_id, vesc_id)
                self._bus.send(can.Message(
                    arbitration_id=arb_id,
                    data=payload,
                    is_extended_id=True,
                ))
            except Exception as exc:
                self.get_logger().debug(
                    f"GET_VALUES send failed (vesc_id={vesc_id}): {exc}"
                )

    def _publish_all(self) -> None:
        with self._state_lock:
            left  = self._state[self._left_id]
            right = self._state[self._right_id]
            # Snapshot dicts while holding lock
            l_dict = self._state_to_dict(left)
            r_dict = self._state_to_dict(right)
            left_alive  = left.is_alive
            right_alive = right.is_alive
            voltage_v   = left.voltage_v if left_alive else right.voltage_v
            combined = {
                "voltage_v":       round(voltage_v, 2),
                "total_current_a": round(left.current_in_a + right.current_in_a, 2),
                "left_rpm":        left.rpm,
                "right_rpm":       right.rpm,
                "left_alive":      left_alive,
                "right_alive":     right_alive,
                "stamp":           time.time(),
            }
            # Build diagnostics while holding lock so values are consistent
            diag_array = self._build_diagnostics(left, right)

        self._pub_left.publish(String(data=json.dumps(l_dict)))
        self._pub_right.publish(String(data=json.dumps(r_dict)))
        self._pub_combined.publish(String(data=json.dumps(combined)))
        self._pub_diag.publish(diag_array)

        # Fault alerting (outside lock — uses snapshots)
        self._check_faults_from_dict(l_dict, self._left_id,  "left")
        self._check_faults_from_dict(r_dict, self._right_id, "right")

    @staticmethod
    def _state_to_dict(s: VescState) -> dict:
        return {
            "can_id":       s.can_id,
            "rpm":          s.rpm,
            "current_a":    round(s.current_a,    2),
            "current_in_a": round(s.current_in_a, 2),
            "duty_cycle":   round(s.duty_cycle,   4),
            "voltage_v":    round(s.voltage_v,    2),
            "temp_fet_c":   round(s.temp_fet_c,   1),
            "temp_motor_c": round(s.temp_motor_c, 1),
            "fault_code":   s.fault_code,
            "fault_name":   s.fault_name,
            "alive":        s.is_alive,
            "stamp":        time.time(),
        }

    def _build_diagnostics(
        self, left: VescState, right: VescState
    ) -> DiagnosticArray:
        array = DiagnosticArray()
        array.header.stamp = self.get_clock().now().to_msg()

        for s, label in ((left, "left"), (right, "right")):
            status = DiagnosticStatus()
            status.name        = f"VESC/{label} (CAN ID {s.can_id})"
            status.hardware_id = f"vesc_can_{s.can_id}"

            if not s.is_alive:
                status.level   = DiagnosticStatus.ERROR
                status.message = "No CAN frames received (offline)"
            elif s.has_fault:
                status.level   = DiagnosticStatus.ERROR
                status.message = f"Fault: {s.fault_name}"
            elif (
                s.temp_fet_c   > self._otf_thresh
                or s.temp_motor_c > self._otm_thresh
                or abs(s.current_a) > self._oc_thresh
            ):
                status.level   = DiagnosticStatus.WARN
                status.message = "Threshold exceeded"
            else:
                status.level   = DiagnosticStatus.OK
                status.message = "OK"

            status.values = [
                KeyValue(key="rpm",          value=str(s.rpm)),
                KeyValue(key="current_a",    value=f"{s.current_a:.2f}"),
                KeyValue(key="current_in_a", value=f"{s.current_in_a:.2f}"),
                KeyValue(key="duty_cycle",   value=f"{s.duty_cycle:.4f}"),
                KeyValue(key="voltage_v",    value=f"{s.voltage_v:.2f}"),
                KeyValue(key="temp_fet_c",   value=f"{s.temp_fet_c:.1f}"),
                KeyValue(key="temp_motor_c", value=f"{s.temp_motor_c:.1f}"),
                KeyValue(key="fault_code",   value=str(s.fault_code)),
                KeyValue(key="fault_name",   value=s.fault_name),
            ]
            array.status.append(status)

        return array

    def _check_faults_from_dict(self, d: dict, vesc_id: int, label: str) -> None:
        """Log fault/threshold warnings; suppress repeated identical messages."""
        if not d["alive"]:
            return

        if abs(d["current_a"]) > self._oc_thresh:
            self.get_logger().warn(
                f"VESC {label}: overcurrent {d['current_a']:.1f} A "
                f"(threshold {self._oc_thresh:.0f} A)"
            )

        if d["temp_fet_c"] > self._otf_thresh:
            self.get_logger().warn(
                f"VESC {label}: FET overtemp {d['temp_fet_c']:.1f} °C "
                f"(threshold {self._otf_thresh:.0f} °C)"
            )

        if d["temp_motor_c"] > self._otm_thresh:
            self.get_logger().warn(
                f"VESC {label}: motor overtemp {d['temp_motor_c']:.1f} °C "
                f"(threshold {self._otm_thresh:.0f} °C)"
            )

        fault_code = d["fault_code"]
        if fault_code != 0 and self._fault_logged.get(vesc_id) != fault_code:
            self.get_logger().error(
                f"VESC {label}: fault {fault_code} ({d['fault_name']})"
            )
            self._fault_logged[vesc_id] = fault_code
        elif fault_code == 0 and self._fault_logged.get(vesc_id, 0) != 0:
            self.get_logger().info(f"VESC {label}: fault cleared")
            self._fault_logged[vesc_id] = 0

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

    def destroy_node(self) -> None:
        self._running = False
        if self._rx_thread is not None:
            self._rx_thread.join(timeout=1.0)
        if self._bus is not None:
            self._bus.shutdown()
        super().destroy_node()


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


if __name__ == "__main__":
    main()
