"""
saltybot_can_node — production FC↔Orin bridge over CAN bus (Issues #680, #672, #685)

In production the FC has NO USB connection to Orin — CAN bus only.
Reads IMU, balance state, barometer, and VESC telemetry from FC over CAN
and publishes them as ROS2 topics.  Sends drive/heartbeat/estop commands
back to FC over CAN.

CAN interface: SocketCAN (CANable USB adapter on vcan0 / can0)

FC → Orin (telemetry):
  0x400  FC_STATUS   int16 pitch_x10, int16 motor_cmd, uint16 vbat_mv,
                     uint8 balance_state, uint8 flags          (10 Hz)
  0x401  FC_VESC     int16 left_rpm_x10, int16 right_rpm_x10,
                     int16 left_cur_x10, int16 right_cur_x10   (10 Hz)
  0x402  FC_IMU      int16 pitch_x10, int16 roll_x10, int16 yaw_x10,
                     uint8 cal_status, uint8 balance_state      (50 Hz)
  0x403  FC_BARO     int32 pressure_pa, int16 temp_x10, int16 alt_cm (1 Hz)

Orin → FC (commands):
  0x300  HEARTBEAT   uint32 sequence counter                    (5 Hz)
  0x301  DRIVE       int16 speed BE, int16 steer BE             (on /cmd_vel)
  0x303  ESTOP       uint8 1=estop, 0=clear
  0x304  LED_CMD     uint8 pattern, uint8 brightness, uint16 duration_ms LE

Published topics:
  /saltybot/imu            (sensor_msgs/Imu)       — from 0x402
  /saltybot/balance_state  (std_msgs/String)        — from 0x402 + 0x400
  /saltybot/barometer      (sensor_msgs/FluidPressure) — from 0x403

Subscribed topics:
  /cmd_vel                 (geometry_msgs/Twist)   — sent as DRIVE
  /saltybot/leds           (std_msgs/String JSON)  — sent as LED_CMD

Parameters:
  can_interface   CAN socket interface name     default: can0
  speed_scale     m/s to ESC units multiplier   default: 1000.0
  steer_scale     rad/s to ESC units multiplier default: -500.0
  heartbeat_hz    HEARTBEAT TX rate (Hz)         default: 5.0
"""

import json
import math
import socket
import struct
import threading
import time

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from geometry_msgs.msg import Twist
from sensor_msgs.msg import Imu, FluidPressure
from std_msgs.msg import String
from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue

# ---- CAN frame IDs ------------------------------------------------

CAN_FC_STATUS   = 0x400
CAN_FC_VESC     = 0x401
CAN_FC_IMU      = 0x402
CAN_FC_BARO     = 0x403

CAN_HEARTBEAT   = 0x300
CAN_DRIVE       = 0x301
CAN_ESTOP       = 0x303
CAN_LED_CMD     = 0x304

# ---- Constants ----------------------------------------------------

IMU_FRAME_ID    = "imu_link"
_STATE_LABEL    = {0: "DISARMED", 1: "ARMED", 2: "TILT_FAULT"}
_CAL_LABEL      = {0: "UNCAL", 1: "GYRO_CAL", 2: "MOUNT_CAL"}

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

def _clamp(v: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, v))


def _can_socket(iface: str) -> socket.socket:
    """Open a raw SocketCAN socket on *iface*."""
    s = socket.socket(socket.AF_CAN, socket.SOCK_RAW, socket.CAN_RAW)
    s.bind((iface,))
    s.settimeout(0.1)
    return s


def _pack_frame(can_id: int, data: bytes) -> bytes:
    """Pack a standard CAN frame for SocketCAN (16-byte struct)."""
    dlc = len(data)
    return struct.pack("=IB3x8s", can_id & 0x1FFFFFFF, dlc,
                       data.ljust(8, b"\x00"))


def _unpack_frame(raw: bytes):
    """Unpack a raw SocketCAN frame; returns (can_id, data_bytes)."""
    can_id, dlc = struct.unpack_from("=IB", raw, 0)
    data = raw[8: 8 + (dlc & 0x0F)]
    return can_id & 0x1FFFFFFF, data


# ---- Node ---------------------------------------------------------

class SaltybotCanNode(Node):
    def __init__(self):
        super().__init__("saltybot_can")

        # ── Parameters ───────────────────────────────────────────────
        self.declare_parameter("can_interface",  "can0")
        self.declare_parameter("speed_scale",    1000.0)
        self.declare_parameter("steer_scale",    -500.0)
        self.declare_parameter("heartbeat_hz",   5.0)

        iface          = self.get_parameter("can_interface").value
        self._speed_sc = self.get_parameter("speed_scale").value
        self._steer_sc = self.get_parameter("steer_scale").value
        hb_hz          = self.get_parameter("heartbeat_hz").value

        # ── QoS ──────────────────────────────────────────────────────
        sensor_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST, depth=10)
        reliable_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST, depth=10)

        # ── Publishers ───────────────────────────────────────────────
        self._imu_pub  = self.create_publisher(Imu,          "/saltybot/imu",           sensor_qos)
        self._bal_pub  = self.create_publisher(String,       "/saltybot/balance_state",  reliable_qos)
        self._baro_pub = self.create_publisher(FluidPressure,"/saltybot/barometer",      reliable_qos)
        self._diag_pub = self.create_publisher(DiagnosticArray, "/diagnostics",           reliable_qos)

        # ── Subscribers ──────────────────────────────────────────────
        self._cmd_sub  = self.create_subscription(
            Twist, "/cmd_vel", self._cmd_vel_cb, reliable_qos)
        self._led_sub  = self.create_subscription(
            String, "/saltybot/leds", self._led_cb, reliable_qos)

        # ── State ────────────────────────────────────────────────────
        self._iface       = iface
        self._sock        = None
        self._sock_lock   = threading.Lock()
        self._hb_seq      = 0
        self._last_speed  = 0
        self._last_steer  = 0
        self._last_state  = -1
        self._last_cal    = -1
        self._last_pitch  = 0.0
        self._last_roll   = 0.0
        self._last_yaw    = 0.0
        self._last_motor  = 0
        self._last_vbat   = 0
        self._last_flags  = 0

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

        # ── Timers ───────────────────────────────────────────────────
        self._rx_timer = self.create_timer(0.01, self._rx_cb)         # 100 Hz poll
        self._hb_timer = self.create_timer(1.0 / hb_hz, self._hb_cb)

        self.get_logger().info(
            f"saltybot_can started — interface={iface} "
            f"(speed_scale={self._speed_sc}, steer_scale={self._steer_sc})"
        )

    # ── CAN socket management ────────────────────────────────────────

    def _open_can(self) -> bool:
        with self._sock_lock:
            try:
                self._sock = _can_socket(self._iface)
                self.get_logger().info(f"CAN socket open: {self._iface}")
                return True
            except OSError as exc:
                self.get_logger().error(f"Cannot open CAN {self._iface}: {exc}")
                self._sock = None
                return False

    def _send(self, can_id: int, data: bytes) -> bool:
        with self._sock_lock:
            if self._sock is None:
                return False
            try:
                self._sock.send(_pack_frame(can_id, data))
                return True
            except OSError as exc:
                self.get_logger().error(
                    f"CAN TX error: {exc}", throttle_duration_sec=2.0)
                return False

    # ── RX poll ──────────────────────────────────────────────────────

    def _rx_cb(self):
        with self._sock_lock:
            if self._sock is None:
                return
            try:
                while True:
                    raw = self._sock.recv(16)
                    can_id, data = _unpack_frame(raw)
                    self._dispatch(can_id, data)
            except socket.timeout:
                pass
            except BlockingIOError:
                pass
            except OSError as exc:
                self.get_logger().error(
                    f"CAN RX error: {exc}", throttle_duration_sec=5.0)
                self._sock = None

        if self._sock is None:
            self._open_can()

    def _dispatch(self, can_id: int, data: bytes):
        now = self.get_clock().now().to_msg()
        if can_id == CAN_FC_IMU and len(data) >= 8:
            self._handle_fc_imu(data, now)
        elif can_id == CAN_FC_STATUS and len(data) >= 8:
            self._handle_fc_status(data)
        elif can_id == CAN_FC_BARO and len(data) >= 8:
            self._handle_fc_baro(data, now)

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

    def _handle_fc_imu(self, data: bytes, stamp):
        pitch_x10, roll_x10, yaw_x10, cal, state = struct.unpack_from("<hhhBB", data, 0)
        pitch = pitch_x10 / 10.0
        roll  = roll_x10  / 10.0
        yaw   = yaw_x10   / 10.0

        self._last_pitch = pitch
        self._last_roll  = roll
        self._last_yaw   = yaw
        self._last_cal   = cal

        # Publish IMU
        imu = Imu()
        imu.header.stamp    = stamp
        imu.header.frame_id = IMU_FRAME_ID
        imu.orientation_covariance[0] = -1.0   # orientation not provided

        # Publish Euler angles via angular_velocity fields (convention matches
        # saltybot_cmd_node — downstream nodes expect this mapping)
        imu.angular_velocity.x = math.radians(pitch)
        imu.angular_velocity.y = math.radians(roll)
        imu.angular_velocity.z = math.radians(yaw)
        cov = math.radians(0.5) ** 2
        imu.angular_velocity_covariance[0] = cov
        imu.angular_velocity_covariance[4] = cov
        imu.angular_velocity_covariance[8] = cov
        imu.linear_acceleration_covariance[0] = -1.0
        self._imu_pub.publish(imu)

        # Publish balance_state JSON
        self._publish_balance_state(pitch, roll, yaw, state, cal, stamp)

        # Log state/cal transitions
        if state != self._last_state:
            self.get_logger().info(
                f"Balance state → {_STATE_LABEL.get(state, f'UNKNOWN({state})')}"
            )
            self._last_state = state
        if cal != self._last_cal:
            self.get_logger().info(
                f"IMU cal status → {_CAL_LABEL.get(cal, f'CAL({cal})')}"
            )

    def _handle_fc_status(self, data: bytes):
        pitch_x10, motor_cmd, vbat_mv, balance_state, flags = \
            struct.unpack_from("<hhhBB", data, 0)
        self._last_motor = motor_cmd
        self._last_vbat  = vbat_mv
        self._last_flags = flags

    def _handle_fc_baro(self, data: bytes, stamp):
        pressure_pa, temp_x10, alt_cm = struct.unpack_from("<ihh", data, 0)
        msg = FluidPressure()
        msg.header.stamp    = stamp
        msg.header.frame_id = "barometer_link"
        msg.fluid_pressure  = float(pressure_pa)      # Pa
        msg.variance        = 0.0
        self._baro_pub.publish(msg)

        diag = DiagnosticArray()
        diag.header.stamp = stamp
        st = DiagnosticStatus()
        st.level       = DiagnosticStatus.OK
        st.name        = "saltybot/barometer"
        st.hardware_id = "bmp280"
        st.message     = "OK"
        st.values = [
            KeyValue(key="pressure_pa", value=str(pressure_pa)),
            KeyValue(key="temp_c",      value=f"{temp_x10 / 10.0:.1f}"),
            KeyValue(key="alt_cm",      value=str(alt_cm)),
        ]
        diag.status.append(st)
        self._diag_pub.publish(diag)

    def _publish_balance_state(self, pitch, roll, yaw, state, cal, stamp):
        state_label = _STATE_LABEL.get(state, f"UNKNOWN({state})")
        cal_label   = _CAL_LABEL.get(cal, f"CAL({cal})")
        payload = {
            "stamp":        f"{stamp.sec}.{stamp.nanosec:09d}",
            "state":        state_label,
            "cal_status":   cal_label,
            "pitch_deg":    round(pitch, 1),
            "roll_deg":     round(roll,  1),
            "yaw_deg":      round(yaw,   1),
            "motor_cmd":    self._last_motor,
            "vbat_mv":      self._last_vbat,
            "jetson_speed": self._last_speed,
            "jetson_steer": self._last_steer,
        }
        str_msg = String()
        str_msg.data = json.dumps(payload)
        self._bal_pub.publish(str_msg)

        diag = DiagnosticArray()
        diag.header.stamp = stamp
        st = DiagnosticStatus()
        st.name        = "saltybot/balance_controller"
<<<<<<< HEAD
        st.hardware_id = "esp32"
=======
        st.hardware_id = "esp32s322"
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        st.message     = state_label
        st.level       = (DiagnosticStatus.OK    if state == 1 else
                          DiagnosticStatus.WARN   if state == 0 else
                          DiagnosticStatus.ERROR)
        st.values = [
            KeyValue(key="pitch_deg",  value=f"{pitch:.1f}"),
            KeyValue(key="roll_deg",   value=f"{roll:.1f}"),
            KeyValue(key="yaw_deg",    value=f"{yaw:.1f}"),
            KeyValue(key="motor_cmd",  value=str(self._last_motor)),
            KeyValue(key="state",      value=state_label),
            KeyValue(key="cal_status", value=cal_label),
        ]
        diag.status.append(st)
        self._diag_pub.publish(diag)

    # ── TX callbacks ─────────────────────────────────────────────────

    def _hb_cb(self):
        """Send HEARTBEAT (0x300) with incrementing sequence counter."""
        data = struct.pack(">I", self._hb_seq & 0xFFFFFFFF)
        self._hb_seq += 1
        self._send(CAN_HEARTBEAT, data)

    def _cmd_vel_cb(self, msg: Twist):
        """Convert Twist → DRIVE (0x301): int16 speed BE, int16 steer BE."""
        speed = int(_clamp(msg.linear.x  * self._speed_sc, -1000.0, 1000.0))
        steer = int(_clamp(msg.angular.z * self._steer_sc, -1000.0, 1000.0))
        self._last_speed = speed
        self._last_steer = steer
        data = struct.pack(">hh", speed, steer)
        self._send(CAN_DRIVE, data)

    def _led_cb(self, msg: String):
        """Parse /saltybot/leds JSON → LED_CMD (0x304).
        Expected JSON: {"pattern": 1, "brightness": 200, "duration_ms": 5000}
        pattern: 0=auto, 1=solid, 2=slow_blink, 3=fast_blink, 4=pulse
        """
        try:
            d = json.loads(msg.data)
        except (ValueError, TypeError) as exc:
            self.get_logger().debug(f"LED JSON parse error: {exc}")
            return
        pattern     = int(d.get("pattern",     0)) & 0xFF
        brightness  = int(d.get("brightness",  255)) & 0xFF
        duration_ms = int(d.get("duration_ms", 0)) & 0xFFFF
        data = struct.pack("<BBH", pattern, brightness, duration_ms)
        self._send(CAN_LED_CMD, data)

    # ── Lifecycle ────────────────────────────────────────────────────

    def destroy_node(self):
        # Send ESTOP on shutdown so FC doesn't keep rolling
        self._send(CAN_ESTOP, bytes([1]))
        with self._sock_lock:
            if self._sock:
                try:
                    self._sock.close()
                except OSError:
                    pass
                self._sock = None
        super().destroy_node()


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


if __name__ == "__main__":
    main()
