#!/usr/bin/env python3
"""
VESC CAN driver node using python-can with SocketCAN.
Subscribes to /cmd_vel (geometry_msgs/Twist) and sends duty cycle
commands to two VESC motor controllers via CAN bus.

Also receives VESC STATUS broadcast frames and publishes telemetry:
  /saltybot/vesc/left   (std_msgs/String JSON)
  /saltybot/vesc/right  (std_msgs/String JSON)

VESC CAN protocol:
  Extended frame ID = (CAN_PACKET_ID << 8) | VESC_ID
  CAN_PACKET_SET_DUTY = 0  -> payload: int32 (duty * 100000)
"""

import json
import struct
import threading
import time
from dataclasses import dataclass, field
from typing import Optional

import can
import rclpy
from geometry_msgs.msg import Twist
from rclpy.node import Node
from std_msgs.msg import String


# ── VESC CAN packet IDs ────────────────────────────────────────────────────────
CAN_PACKET_SET_DUTY   = 0
CAN_PACKET_SET_RPM    = 3
CAN_PACKET_STATUS     = 9    # RPM, phase current, duty
CAN_PACKET_STATUS_4   = 16   # FET temp, motor temp, input current
CAN_PACKET_STATUS_5   = 27   # Input voltage

FAULT_NAMES = {
    0: "NONE",                        1: "OVER_VOLTAGE",
    2: "UNDER_VOLTAGE",               3: "DRV",
    4: "ABS_OVER_CURRENT",            5: "OVER_TEMP_FET",
    6: "OVER_TEMP_MOTOR",             7: "GATE_DRIVER_OVER_VOLTAGE",
    8: "GATE_DRIVER_UNDER_VOLTAGE_3V3", 9: "MCU_UNDER_VOLTAGE",
    10: "BOOTING_FROM_WATCHDOG_RESET", 11: "ENCODER_SPI_FAULT",
}

ALIVE_TIMEOUT_S = 1.0


@dataclass
class VescState:
    can_id:       int
    rpm:          int   = 0
    current_a:    float = 0.0
    current_in_a: float = 0.0
    duty_cycle:   float = 0.0
    voltage_v:    float = 0.0
    temp_fet_c:   float = 0.0
    temp_motor_c: float = 0.0
    fault_code:   int   = 0
    last_ts:      float = field(default_factory=lambda: 0.0)

    @property
    def is_alive(self) -> bool:
        return (time.monotonic() - self.last_ts) < ALIVE_TIMEOUT_S

    @property
    def fault_name(self) -> str:
        return FAULT_NAMES.get(self.fault_code, f"UNKNOWN_{self.fault_code}")

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


def make_can_id(packet_id: int, vesc_id: int) -> int:
    """Build extended CAN frame ID for VESC."""
    return (packet_id << 8) | vesc_id


class VescCanDriver(Node):
    def __init__(self):
        super().__init__('vesc_can_driver')

        # Declare parameters
        self.declare_parameter('interface', 'can0')
        self.declare_parameter('left_motor_id', 56)
        self.declare_parameter('right_motor_id', 68)
        self.declare_parameter('wheel_separation', 0.5)
        self.declare_parameter('wheel_radius', 0.1)
        self.declare_parameter('max_speed', 5.0)
        self.declare_parameter('command_timeout', 1.0)
        self.declare_parameter('telemetry_rate_hz', 10)

        # Read parameters
        self.interface   = self.get_parameter('interface').value
        self.left_id     = self.get_parameter('left_motor_id').value
        self.right_id    = self.get_parameter('right_motor_id').value
        self.wheel_sep   = self.get_parameter('wheel_separation').value
        self.max_speed   = self.get_parameter('max_speed').value
        self.cmd_timeout = self.get_parameter('command_timeout').value
        tel_hz           = int(self.get_parameter('telemetry_rate_hz').value)

        # Open CAN bus
        try:
            self.bus = can.interface.Bus(
                channel=self.interface,
                bustype='socketcan',
            )
            self.get_logger().info(f'CAN bus opened: {self.interface}')
        except Exception as e:
            self.get_logger().error(f'Failed to open CAN bus: {e}')
            raise

        self._last_cmd_time = time.monotonic()

        # Per-VESC telemetry state
        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()

        # Telemetry publishers
        self._pub_left  = self.create_publisher(String, '/saltybot/vesc/left',  10)
        self._pub_right = self.create_publisher(String, '/saltybot/vesc/right', 10)

        # CAN RX background thread
        self._running   = True
        self._rx_thread = threading.Thread(target=self._rx_loop, name='vesc_rx', daemon=True)
        self._rx_thread.start()

        # Subscriber
        self.create_subscription(Twist, '/cmd_vel_smoothed', self._cmd_vel_cb, 10)

        # Watchdog + telemetry publish timer
        self.create_timer(1.0 / max(1, tel_hz), self._watchdog_and_publish_cb)

        self.get_logger().info(
            f'VESC CAN driver ready — left={self.left_id} right={self.right_id} '
            f'telemetry={tel_hz} Hz'
        )

    # ------------------------------------------------------------------
    # Command velocity callback
    # ------------------------------------------------------------------

    def _cmd_vel_cb(self, msg: Twist):
        self._last_cmd_time = time.monotonic()
        linear  = msg.linear.x
        angular = msg.angular.z

        left_speed  = linear - (angular * self.wheel_sep / 2.0)
        right_speed = linear + (angular * self.wheel_sep / 2.0)

        left_duty  = max(-1.0, min(1.0, left_speed  / self.max_speed))
        right_duty = max(-1.0, min(1.0, right_speed / self.max_speed))

        self._send_duty(self.left_id,  left_duty)
        self._send_duty(self.right_id, right_duty)

    # ------------------------------------------------------------------
    # Watchdog + telemetry publish (combined timer callback)
    # ------------------------------------------------------------------

    def _watchdog_and_publish_cb(self):
        elapsed = time.monotonic() - self._last_cmd_time
        if elapsed > self.cmd_timeout:
            self._send_duty(self.left_id,  0.0)
            self._send_duty(self.right_id, 0.0)

        with self._state_lock:
            l_dict = self._state[self.left_id].to_dict()
            r_dict = self._state[self.right_id].to_dict()

        self._pub_left.publish(String(data=json.dumps(l_dict)))
        self._pub_right.publish(String(data=json.dumps(r_dict)))

    # ------------------------------------------------------------------
    # CAN RX background thread
    # ------------------------------------------------------------------

    def _rx_loop(self) -> None:
        """Receive VESC STATUS broadcast frames and update telemetry state."""
        while self._running:
            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:
        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 and len(data) >= 8:
                rpm          = struct.unpack('>i', data[0:4])[0]
                current_a    = struct.unpack('>h', data[4:6])[0] / 10.0
                duty_cycle   = struct.unpack('>h', data[6:8])[0] / 1000.0
                s.rpm        = rpm
                s.current_a  = current_a
                s.duty_cycle = duty_cycle
                s.last_ts    = now

            elif packet_type == CAN_PACKET_STATUS_4 and len(data) >= 6:
                s.temp_fet_c   = struct.unpack('>h', data[0:2])[0] / 10.0
                s.temp_motor_c = struct.unpack('>h', data[2:4])[0] / 10.0
                s.current_in_a = struct.unpack('>h', data[4:6])[0] / 10.0

            elif packet_type == CAN_PACKET_STATUS_5 and len(data) >= 2:
                s.voltage_v = struct.unpack('>h', data[0:2])[0] / 10.0

    # ------------------------------------------------------------------
    # CAN send helpers
    # ------------------------------------------------------------------

    def _send_duty(self, vesc_id: int, duty: float):
        """Send CAN_PACKET_SET_DUTY to a VESC controller."""
        can_id  = make_can_id(CAN_PACKET_SET_DUTY, vesc_id)
        payload = struct.pack('>i', int(duty * 100000))
        msg = can.Message(
            arbitration_id=can_id,
            data=payload,
            is_extended_id=True,
        )
        try:
            self.bus.send(msg)
        except can.CanError as e:
            self.get_logger().error(f'CAN send error (id={vesc_id}): {e}')

    def destroy_node(self):
        self._running = False
        self._rx_thread.join(timeout=1.0)
        self._send_duty(self.left_id,  0.0)
        self._send_duty(self.right_id, 0.0)
        self.bus.shutdown()
        super().destroy_node()


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


if __name__ == '__main__':
    main()
