#!/usr/bin/env python3
"""
can_bridge_node.py — ROS2 node bridging the SaltyBot Orin to the ESP32-S3 BALANCE motor
controller and VESC motor controllers over CAN bus.

Spec: docs/SAUL-TEE-SYSTEM-REFERENCE.md §6 (2026-04-04)

Subscriptions
-------------
  /cmd_vel          geometry_msgs/Twist          → VESC speed commands (CAN)
  /estop            std_msgs/Bool                → ESP32-S3 BALANCE e-stop (CAN)

Publications
------------
  /can/imu                sensor_msgs/Imu              ESP32-S3 BALANCE IMU telemetry
  /can/battery            sensor_msgs/BatteryState     ESP32-S3 BALANCE battery telemetry
  /can/vesc/left/state    std_msgs/Float32MultiArray   Left VESC state
  /can/vesc/right/state   std_msgs/Float32MultiArray   Right VESC state
  /can/connection_status  std_msgs/String              "connected" | "disconnected"

Parameters
----------
  can_interface     str    CAN socket name (default: slcan0)
  speed_scale       float  /cmd_vel linear.x (m/s) → motor units  (default: 1000.0)
  steer_scale       float  /cmd_vel angular.z (rad/s) → motor units (default: -500.0)
  command_timeout_s float  watchdog zero-vel threshold (default: 0.5)
"""

import json
import threading
import time
from typing import Optional

import can
import rclpy
from geometry_msgs.msg import Twist
from rclpy.node import Node
from sensor_msgs.msg import BatteryState
from std_msgs.msg import Bool, Float32MultiArray, String

from saltybot_can_bridge.balance_protocol import (
    MAMBA_CMD_ESTOP,
    MAMBA_CMD_MODE,
    MAMBA_CMD_VELOCITY,
    MAMBA_TELEM_BATTERY,
    MAMBA_TELEM_IMU,
    VESC_TELEM_STATE,
    ORIN_CAN_ID_FC_PID_ACK,
    ORIN_CAN_ID_PID_SET,
    MODE_DRIVE,
    MODE_IDLE,
    encode_drive_cmd,
    encode_arm_cmd,
    encode_estop_cmd,
    decode_attitude,
    decode_battery,
    decode_vesc_status1,
)

# Reconnect attempt interval when CAN bus is lost
_RECONNECT_INTERVAL_S: float = 5.0

# Watchdog tick rate (Hz); sends zero DRIVE when /cmd_vel is silent
_WATCHDOG_HZ: float = 10.0


class CanBridgeNode(Node):
    """CAN bus bridge between Orin ROS2 and ESP32-S3 BALANCE / VESC controllers."""

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

        # ── Parameters ────────────────────────────────────────────────────
        self.declare_parameter("can_interface",     "slcan0")
        self.declare_parameter("left_vesc_can_id",  VESC_LEFT_ID)
        self.declare_parameter("right_vesc_can_id", VESC_RIGHT_ID)
        self.declare_parameter("speed_scale",       1000.0)
        self.declare_parameter("steer_scale",       -500.0)
        self.declare_parameter("command_timeout_s", 0.5)

        self._iface         = self.get_parameter("can_interface").value
        self._left_vesc_id  = self.get_parameter("left_vesc_can_id").value
        self._right_vesc_id = self.get_parameter("right_vesc_can_id").value
        self._speed_scale   = self.get_parameter("speed_scale").value
        self._steer_scale   = self.get_parameter("steer_scale").value
        self._cmd_timeout   = self.get_parameter("command_timeout_s").value

        # ── State ─────────────────────────────────────────────────────────
        self._bus: Optional[can.BusABC] = None
        self._connected: bool = False
        self._last_cmd_time: float = time.monotonic()
        self._lock = threading.Lock()

        # ── Publishers ────────────────────────────────────────────────────
        self._pub_attitude  = self.create_publisher(String,           "/saltybot/attitude",      10)
        self._pub_balance   = self.create_publisher(String,           "/saltybot/balance_state", 10)
        self._pub_battery   = self.create_publisher(BatteryState,     "/can/battery",            10)
        self._pub_vesc_left = self.create_publisher(Float32MultiArray,"/can/vesc/left/state",    10)
        self._pub_vesc_right= self.create_publisher(Float32MultiArray,"/can/vesc/right/state",   10)
        self._pub_status    = self.create_publisher(String,           "/can/connection_status",  10)

        # ── Subscriptions ─────────────────────────────────────────────────
        self.create_subscription(Twist, "/cmd_vel",       self._cmd_vel_cb, 10)
        self.create_subscription(Bool,  "/estop",         self._estop_cb,   10)
        self.create_subscription(Bool,  "/saltybot/arm",  self._arm_cb,     10)

        # ── Timers ────────────────────────────────────────────────────────
        self.create_timer(1.0 / _WATCHDOG_HZ,  self._watchdog_cb)
        self.create_timer(_RECONNECT_INTERVAL_S, self._reconnect_cb)

        # ── Open CAN ──────────────────────────────────────────────────────
        self._try_connect()

        # ── Background reader thread ──────────────────────────────────────
        self._reader_thread = threading.Thread(
            target=self._reader_loop, daemon=True, name="can_reader"
        )
        self._reader_thread.start()

        self.get_logger().info(
            f"can_bridge_node ready — iface={self._iface} "
            f"left_vesc={self._left_vesc_id} right_vesc={self._right_vesc_id} "
            f"speed_scale={self._speed_scale} steer_scale={self._steer_scale}"
        )

    # ── Connection management ──────────────────────────────────────────────

    def _try_connect(self) -> None:
        with self._lock:
            if self._connected:
                return
            try:
                self._bus = can.interface.Bus(channel=self._iface, bustype="socketcan")
                self._connected = True
                self.get_logger().info(f"CAN bus connected: {self._iface}")
                self._publish_status("connected")
            except Exception as exc:
                self.get_logger().warning(
                    f"CAN bus not available ({self._iface}): {exc} "
                    f"— will retry every {_RECONNECT_INTERVAL_S:.0f} s"
                )
                self._connected = False
                self._publish_status("disconnected")

    def _reconnect_cb(self) -> None:
        if not self._connected:
            self._try_connect()

    def _handle_can_error(self, exc: Exception, context: str) -> None:
        self.get_logger().warning(f"CAN error in {context}: {exc}")
        with self._lock:
            if self._bus is not None:
                try:
                    self._bus.shutdown()
                except Exception:
                    pass
            self._bus = None
            self._connected = False
        self._publish_status("disconnected")

    # ── ROS callbacks ─────────────────────────────────────────────────────

    def _cmd_vel_cb(self, msg: Twist) -> None:
        """Convert /cmd_vel Twist to ORIN_CMD_DRIVE over CAN."""
        self._last_cmd_time = time.monotonic()
        if not self._connected:
            return

        # Differential drive decomposition — individual wheel speeds in m/s.
        # The VESC nodes interpret linear velocity directly; angular is handled
        # by the sign difference between left and right.
        linear = msg.linear.x
        angular = msg.angular.z

        # Forward left = forward right for pure translation; for rotation
        # left slows and right speeds up (positive angular = CCW = left turn).
        # The ESP32-S3 BALANCE velocity command carries both wheels independently.
        left_mps = linear - angular
        right_mps = linear + angular

        payload = encode_velocity_cmd(left_mps, right_mps)
        self._send_can(MAMBA_CMD_VELOCITY, payload, "cmd_vel")

        # Keep ESP32-S3 BALANCE in DRIVE mode while receiving commands
        self._send_can(MAMBA_CMD_MODE, encode_mode_cmd(MODE_DRIVE), "cmd_vel mode")

    def _estop_cb(self, msg: Bool) -> None:
        """Forward /estop to ESP32-S3 BALANCE over CAN."""
        if not self._connected:
            return
        if msg.data:
            self._send_can(
                MAMBA_CMD_MODE, encode_mode_cmd(MODE_ESTOP), "estop mode"
            )
            self.get_logger().warning("E-stop asserted — sent ESTOP to ESP32-S3 BALANCE")

    # ── Watchdog ──────────────────────────────────────────────────────────

    def _watchdog_cb(self) -> None:
        """If /cmd_vel is silent for command_timeout_s, send zero DRIVE (acts as keepalive)."""
        if not self._connected:
            return
        if time.monotonic() - self._last_cmd_time > self._cmd_timeout:
            self._send_can(ORIN_CMD_DRIVE, encode_drive_cmd(0, 0, MODE_IDLE), "watchdog")

    # ── CAN send helper ───────────────────────────────────────────────────

    def _send_can(self, arb_id: int, data: bytes, context: str,
                  extended: bool = False) -> None:
        with self._lock:
            if not self._connected or self._bus is None:
                return
            bus = self._bus
        msg = can.Message(arbitration_id=arb_id, data=data,
                          is_extended_id=extended)
        try:
            bus.send(msg, timeout=0.05)
        except can.CanError as exc:
            self._handle_can_error(exc, f"send({context})")

    # ── Background CAN reader ─────────────────────────────────────────────

    def _reader_loop(self) -> None:
        while rclpy.ok():
            with self._lock:
                connected, bus = self._connected, self._bus
            if not connected or bus is None:
                time.sleep(0.1)
                continue
            try:
                frame = bus.recv(timeout=0.5)
            except can.CanError as exc:
                self._handle_can_error(exc, "reader_loop recv")
                continue
            if frame is None:
                continue
            self._dispatch_frame(frame)

    def _dispatch_frame(self, frame: can.Message) -> None:
        arb_id = frame.arbitration_id
        data   = bytes(frame.data)
        vesc_l = (VESC_STATUS_1 << 8) | self._left_vesc_id
        vesc_r = (VESC_STATUS_1 << 8) | self._right_vesc_id
        try:
            if arb_id == ESP32_TELEM_ATTITUDE:
                self._handle_attitude(data)
            elif arb_id == ESP32_TELEM_BATTERY:
                self._handle_battery(data)
            elif arb_id == vesc_l:
                t = decode_vesc_status1(self._left_vesc_id, data)
                m = Float32MultiArray()
                m.data = [t.erpm, t.duty, 0.0, t.current]
                self._pub_vesc_left.publish(m)
            elif arb_id == vesc_r:
                t = decode_vesc_status1(self._right_vesc_id, data)
                m = Float32MultiArray()
                m.data = [t.erpm, t.duty, 0.0, t.current]
                self._pub_vesc_right.publish(m)
        except Exception as exc:
            self.get_logger().warning(
                f"Error parsing CAN frame 0x{arb_id:03X}: {exc}"
            )

    # ── Frame handlers ────────────────────────────────────────────────────

    _STATE_LABEL = {0: "IDLE", 1: "RUNNING", 2: "FAULT"}

    def _handle_attitude(self, data: bytes) -> None:
        """ATTITUDE (0x400): pitch, speed, yaw_rate, state, flags → /saltybot/attitude."""
        t = decode_attitude(data)
        now = self.get_clock().now().to_msg()
        payload = {
            "pitch_deg":  round(t.pitch_deg, 2),
            "speed_mps":  round(t.speed, 3),
            "yaw_rate":   round(t.yaw_rate, 3),
            "state":      t.state,
            "state_label": self._STATE_LABEL.get(t.state, f"UNKNOWN({t.state})"),
            "flags":      t.flags,
            "ts":         f"{now.sec}.{now.nanosec:09d}",
        }
        msg = String()
        msg.data = json.dumps(payload)
        self._pub_attitude.publish(msg)
        self._pub_balance.publish(msg)   # keep /saltybot/balance_state alive

    def _handle_battery(self, data: bytes) -> None:
        """BATTERY (0x401): vbat_mv, fault_code, rssi → /can/battery."""
        t = decode_battery(data)
        msg = BatteryState()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.voltage = t.vbat_mv / 1000.0
        msg.present = True
        msg.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_DISCHARGING
        self._pub_battery.publish(msg)

    # ── Status helper ─────────────────────────────────────────────────────

    def _publish_status(self, status: str) -> None:
        msg = String()
        msg.data = status
        self._pub_status.publish(msg)

    # ── Shutdown ──────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        if self._connected and self._bus is not None:
            try:
                self._send_can(ORIN_CMD_DRIVE, encode_drive_cmd(0, 0, MODE_IDLE), "shutdown")
                self._send_can(ORIN_CMD_ARM,   encode_arm_cmd(False), "shutdown")
            except Exception:
                pass
            try:
                self._bus.shutdown()
            except Exception:
                pass
        super().destroy_node()


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


if __name__ == "__main__":
    main()
