"""
<<<<<<< HEAD
cmd_vel_bridge_node — Nav2 /cmd_vel → ESP32 BALANCE drive command bridge.
=======
cmd_vel_bridge_node — Nav2 /cmd_vel → ESP32-S3 drive command bridge.
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

Extends the basic saltybot_cmd_node with four additions required for safe
autonomous operation on a self-balancing robot:

  1. Velocity limits   — clamp linear.x / angular.z before scaling so Nav2
                         planner misconfiguration can't command unsafe speed.
  2. Smooth ramp       — interpolate ESC commands toward target at ramp_rate
                         ESC-units/s to prevent the sudden torque changes that
                         topple a balance robot.
  3. Deadman switch    — if /cmd_vel is silent for cmd_vel_timeout seconds,
                         zero targets immediately (Nav2 node crash / planner
                         stall → robot coasts to stop rather than running away).
<<<<<<< HEAD
  4. Mode gate         — only issue non-zero drive commands when ESP32 BALANCE reports
=======
  4. Mode gate         — only issue non-zero drive commands when ESP32-S3 reports
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
                         md=2 (AUTONOMOUS). In any other mode (RC_MANUAL,
                         RC_ASSISTED) Jetson cannot override the RC pilot.
                         On mode re-entry current ramp state resets to 0 so
                         acceleration is always smooth from rest.

Serial protocol (C<speed>,<steer>\\n / H\\n — same as saltybot_cmd_node):
  C<spd>,<str>\\n  — drive command. speed/steer: -1000..+1000 integers.
<<<<<<< HEAD
  H\\n             — heartbeat. ESP32 BALANCE reverts steer to 0 after 500ms silence.

Telemetry (50 Hz from ESP32 BALANCE):
=======
  H\\n             — heartbeat. ESP32-S3 reverts steer to 0 after 500ms silence.

Telemetry (50 Hz from ESP32-S3):
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
  Same RX/publish pipeline as saltybot_cmd_node.
  The "md" field (0=MANUAL,1=ASSISTED,2=AUTO) is parsed for the mode gate.

Topics published:
  /saltybot/imu            sensor_msgs/Imu
  /saltybot/balance_state  std_msgs/String (JSON)
  /saltybot/cmd            std_msgs/String — formatted command for observability
  /diagnostics             diagnostic_msgs/DiagnosticArray

Topics subscribed:
  /cmd_vel                 geometry_msgs/Twist

This node owns the serial port exclusively. Do NOT run alongside
serial_bridge_node or saltybot_cmd_node on the same port.
"""

import json
import math
import threading
import time

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

import serial

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

_STATE_LABEL = {0: "DISARMED", 1: "ARMED", 2: "TILT_FAULT"}
_MODE_LABEL  = {0: "RC_MANUAL", 1: "RC_ASSISTED", 2: "AUTONOMOUS"}
IMU_FRAME_ID = "imu_link"
MODE_AUTONOMOUS = 2


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


def _ramp_toward(current: int, target: int, max_step: int) -> int:
    """Step `current` toward `target` by at most `max_step`."""
    diff = target - current
    if abs(diff) <= max_step:
        return target
    return current + (max_step if diff > 0 else -max_step)


class CmdVelBridgeNode(Node):
    def __init__(self):
        super().__init__("cmd_vel_bridge")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("serial_port",      "/dev/ttyACM0")
        self.declare_parameter("baud_rate",        921600)
        self.declare_parameter("timeout",          0.05)
        self.declare_parameter("reconnect_delay",  2.0)
        self.declare_parameter("heartbeat_period", 0.2)
        # Velocity limits (applied before scale → hard cap on robot speed)
        self.declare_parameter("max_linear_vel",   0.5)    # m/s
        self.declare_parameter("max_angular_vel",  2.0)    # rad/s
        # Scaling: physical units → ESC counts (-1000..+1000)
        self.declare_parameter("speed_scale",      1000.0) # m/s  → ESC
        self.declare_parameter("steer_scale",      -500.0) # rad/s→ ESC (neg=flip CCW→right)
        # Ramp: max ESC-unit change per second (both speed and steer)
        self.declare_parameter("ramp_rate",        500)    # ESC units/s
        # Deadman: zero targets if /cmd_vel silent longer than this
        self.declare_parameter("cmd_vel_timeout",  0.5)    # seconds

        port     = self.get_parameter("serial_port").value
        baud     = self.get_parameter("baud_rate").value
        timeout  = self.get_parameter("timeout").value
        self._reconnect_delay  = self.get_parameter("reconnect_delay").value
        self._hb_period        = self.get_parameter("heartbeat_period").value
        self._max_linear_vel   = self.get_parameter("max_linear_vel").value
        self._max_angular_vel  = self.get_parameter("max_angular_vel").value
        self._speed_scale      = self.get_parameter("speed_scale").value
        self._steer_scale      = self.get_parameter("steer_scale").value
        self._ramp_rate        = int(self.get_parameter("ramp_rate").value)
        self._cmd_vel_timeout  = self.get_parameter("cmd_vel_timeout").value

        # Control loop runs at 50 Hz; ramp step per tick = ramp_rate / 50
        _CONTROL_HZ = 50
        self._ramp_step = max(1, self._ramp_rate // _CONTROL_HZ)

        # ── 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._balance_pub = self.create_publisher(String, "/saltybot/balance_state",  reliable_qos)
        self._diag_pub    = self.create_publisher(DiagnosticArray, "/diagnostics",    reliable_qos)
        self._cmd_pub     = self.create_publisher(String, "/saltybot/cmd",            reliable_qos)

        # ── Subscription ──────────────────────────────────────────────────────
        self._cmd_vel_sub = self.create_subscription(
            Twist, "/cmd_vel", self._cmd_vel_cb, reliable_qos)

        # ── Drive state ───────────────────────────────────────────────────────
        self._target_speed  = 0   # desired ESC units from latest /cmd_vel
        self._target_steer  = 0
        self._current_speed = 0   # ramped output actually sent
        self._current_steer = 0
        self._last_cmd_vel  = 0.0 # wall clock (seconds) of last /cmd_vel msg
        self._esp32_mode    = 0   # parsed "md" field: 0=MANUAL,1=ASSISTED,2=AUTO
        self._last_state    = -1
        self._frame_count   = 0
        self._error_count   = 0
        self._last_speed_sent = 0
        self._last_steer_sent = 0

        # ── Serial ────────────────────────────────────────────────────────────
        self._port_name = port
        self._baud      = baud
        self._timeout   = timeout
        self._ser: serial.Serial | None = None
        self._ser_lock  = threading.Lock()
        self._open_serial()

        # ── Timers ────────────────────────────────────────────────────────────
<<<<<<< HEAD
        # Telemetry read at 100 Hz (ESP32 BALANCE sends at 50 Hz)
=======
        # Telemetry read at 100 Hz (ESP32-S3 sends at 50 Hz)
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        self._read_timer    = self.create_timer(0.01,             self._read_cb)
        # Control loop at 50 Hz: ramp + deadman + mode gate + send
        self._control_timer = self.create_timer(1.0 / _CONTROL_HZ, self._control_cb)
        # Heartbeat TX
        self._hb_timer      = self.create_timer(self._hb_period,  self._heartbeat_cb)

        self.get_logger().info(
            f"cmd_vel_bridge started — {port} @ {baud} baud | "
            f"max_vel={self._max_linear_vel:.1f}m/s "
            f"max_ang={self._max_angular_vel:.1f}rad/s "
            f"ramp={self._ramp_rate}ESC/s "
            f"deadman={self._cmd_vel_timeout:.1f}s"
        )

    # ── Serial management ─────────────────────────────────────────────────────

    def _open_serial(self) -> bool:
        with self._ser_lock:
            try:
                self._ser = serial.Serial(
                    port=self._port_name, baudrate=self._baud,
                    timeout=self._timeout)
                self._ser.reset_input_buffer()
                self.get_logger().info(f"Serial open: {self._port_name}")
                return True
            except serial.SerialException as exc:
                self.get_logger().error(f"Cannot open {self._port_name}: {exc}")
                self._ser = None
                return False

    def _close_serial(self):
        with self._ser_lock:
            if self._ser and self._ser.is_open:
                try:
                    self._ser.close()
                except Exception:
                    pass
            self._ser = None

    def _write(self, data: bytes) -> bool:
        with self._ser_lock:
            if self._ser is None or not self._ser.is_open:
                return False
            try:
                self._ser.write(data)
                return True
            except serial.SerialException as exc:
                self.get_logger().error(f"Serial write error: {exc}")
                self._ser = None
                return False

    # ── /cmd_vel callback ─────────────────────────────────────────────────────

    def _cmd_vel_cb(self, msg: Twist):
        """Clamp velocity to limits then convert to ESC units."""
        lin_x = _clamp(msg.linear.x,  -self._max_linear_vel,  self._max_linear_vel)
        ang_z = _clamp(msg.angular.z, -self._max_angular_vel, self._max_angular_vel)
        self._target_speed = int(_clamp(lin_x * self._speed_scale, -1000.0, 1000.0))
        self._target_steer = int(_clamp(ang_z * self._steer_scale, -1000.0, 1000.0))
        self._last_cmd_vel = time.monotonic()

    # ── 50 Hz control loop ────────────────────────────────────────────────────

    def _control_cb(self):
        now = time.monotonic()

        # Deadman: zero targets if /cmd_vel went silent
        if (self._last_cmd_vel > 0.0 and
                now - self._last_cmd_vel > self._cmd_vel_timeout):
            self._target_speed = 0
            self._target_steer = 0

        # Mode gate: in non-AUTONOMOUS mode, zero and reset ramp state so
        # re-entry always accelerates smoothly from 0.
        if self._esp32_mode != MODE_AUTONOMOUS:
            self._current_speed = 0
            self._current_steer = 0
            speed, steer = 0, 0
        else:
            # Smooth ramp toward targets
            self._current_speed = _ramp_toward(
                self._current_speed, self._target_speed, self._ramp_step)
            self._current_steer = _ramp_toward(
                self._current_steer, self._target_steer, self._ramp_step)
            speed = self._current_speed
            steer = self._current_steer

<<<<<<< HEAD
        # Send to ESP32 BALANCE
=======
        # Send to ESP32-S3
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        frame = f"C{speed},{steer}\n".encode("ascii")
        if not self._write(frame):
            self.get_logger().warn(
                "Cannot send cmd — serial not open",
                throttle_duration_sec=2.0)

        self._last_speed_sent = speed
        self._last_steer_sent = steer

        # Publish command for observability / rosbag
        cmd_msg = String()
        cmd_msg.data = f"C{speed},{steer}"
        self._cmd_pub.publish(cmd_msg)

    # ── Heartbeat TX ──────────────────────────────────────────────────────────

    def _heartbeat_cb(self):
<<<<<<< HEAD
        """H\\n keeps ESP32 BALANCE jetson_cmd heartbeat alive regardless of mode."""
=======
        """H\\n keeps ESP32-S3 jetson_cmd heartbeat alive regardless of mode."""
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        self._write(b"H\n")

    # ── Telemetry RX ──────────────────────────────────────────────────────────

    def _read_cb(self):
        with self._ser_lock:
            if self._ser is None or not self._ser.is_open:
                pass
            else:
                try:
                    lines = []
                    while self._ser.in_waiting:
                        raw = self._ser.readline()
                        if raw:
                            lines.append(raw)
                except serial.SerialException as exc:
                    self.get_logger().error(f"Serial read error: {exc}")
                    self._ser = None
                    lines = []

        for raw in lines:
            self._parse_and_publish(raw)

        with self._ser_lock:
            if self._ser is None:
                self.get_logger().warn(
                    "Serial lost — reconnecting…",
                    throttle_duration_sec=self._reconnect_delay)
        if self._ser is None:
            self._open_serial()

    def _parse_and_publish(self, raw: bytes):
        try:
            text = raw.decode("ascii", errors="ignore").strip()
            if not text:
                return
            data = json.loads(text)
        except (ValueError, UnicodeDecodeError) as exc:
            self.get_logger().debug(f"Parse error ({exc}): {raw!r}")
            self._error_count += 1
            return

        now = self.get_clock().now().to_msg()

        if "err" in data:
            self._publish_imu_fault(data["err"], now)
            return

        required = ("p", "r", "e", "ig", "m", "s", "y")
        if not all(k in data for k in required):
            self.get_logger().debug(f"Incomplete frame: {text}")
            return

        pitch_deg = data["p"] / 10.0
        roll_deg  = data["r"] / 10.0
        yaw_deg   = data["y"] / 10.0
        error_deg = data["e"] / 10.0
        integral  = data["ig"] / 10.0
        motor_cmd = float(data["m"])
        state     = int(data["s"])
        mode      = int(data.get("md", 0))  # 0=MANUAL if not present

        self._esp32_mode = mode
        self._frame_count += 1

        self._publish_imu(pitch_deg, roll_deg, yaw_deg, now)
        self._publish_balance_state(
            pitch_deg, roll_deg, yaw_deg,
            error_deg, integral, motor_cmd, state, mode, now)

        if state != self._last_state:
            self.get_logger().info(
                f"Balance state → {_STATE_LABEL.get(state, f'UNKNOWN({state})')}"
                f"  mode={_MODE_LABEL.get(mode, str(mode))}"
            )
            self._last_state = state

    def _publish_imu(self, pitch_deg, roll_deg, yaw_deg, stamp):
        msg = Imu()
        msg.header.stamp    = stamp
        msg.header.frame_id = IMU_FRAME_ID
        msg.orientation_covariance[0] = -1.0
        msg.angular_velocity.x = math.radians(pitch_deg)
        msg.angular_velocity.y = math.radians(roll_deg)
        msg.angular_velocity.z = math.radians(yaw_deg)
        cov = math.radians(0.5) ** 2
        msg.angular_velocity_covariance[0] = cov
        msg.angular_velocity_covariance[4] = cov
        msg.angular_velocity_covariance[8] = cov
        msg.linear_acceleration_covariance[0] = -1.0
        self._imu_pub.publish(msg)

    def _publish_balance_state(
        self, pitch, roll, yaw, error, integral,
        motor_cmd, state, mode, stamp
    ):
        state_label = _STATE_LABEL.get(state, f"UNKNOWN({state})")
        mode_label  = _MODE_LABEL.get(mode, str(mode))
        payload = {
            "stamp":         f"{stamp.sec}.{stamp.nanosec:09d}",
            "state":         state_label,
            "mode":          mode_label,
            "pitch_deg":     round(pitch, 1),
            "roll_deg":      round(roll, 1),
            "yaw_deg":       round(yaw, 1),
            "pid_error_deg": round(error, 1),
            "integral":      round(integral, 1),
            "motor_cmd":     int(motor_cmd),
            "bridge_speed":  self._last_speed_sent,
            "bridge_steer":  self._last_steer_sent,
            "frames":        self._frame_count,
            "parse_errors":  self._error_count,
        }
        str_msg = String()
        str_msg.data = json.dumps(payload)
        self._balance_pub.publish(str_msg)

        diag = DiagnosticArray()
        diag.header.stamp = stamp
        status = DiagnosticStatus()
        status.name        = "saltybot/balance_controller"
<<<<<<< HEAD
        status.hardware_id = "esp32"
=======
        status.hardware_id = "esp32s322"
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        status.message     = f"{state_label} [{mode_label}]"
        status.level = (
            DiagnosticStatus.OK    if state == 1 else
            DiagnosticStatus.WARN  if state == 0 else
            DiagnosticStatus.ERROR
        )
        status.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="pid_error",    value=f"{error:.1f}"),
            KeyValue(key="integral",     value=f"{integral:.1f}"),
            KeyValue(key="motor_cmd",    value=str(int(motor_cmd))),
            KeyValue(key="mode",         value=mode_label),
            KeyValue(key="bridge_speed", value=str(self._last_speed_sent)),
            KeyValue(key="bridge_steer", value=str(self._last_steer_sent)),
            KeyValue(key="state",        value=state_label),
        ]
        diag.status.append(status)
        self._diag_pub.publish(diag)

    def _publish_imu_fault(self, errno: int, stamp):
        diag = DiagnosticArray()
        diag.header.stamp = stamp
        status = DiagnosticStatus()
        status.level       = DiagnosticStatus.ERROR
        status.name        = "saltybot/balance_controller"
<<<<<<< HEAD
        status.hardware_id = "esp32"
        status.message     = f"IMU fault errno={errno}"
        diag.status.append(status)
        self._diag_pub.publish(diag)
        self.get_logger().error(f"ESP32 BALANCE IMU fault: errno={errno}")
=======
        status.hardware_id = "esp32s322"
        status.message     = f"IMU fault errno={errno}"
        diag.status.append(status)
        self._diag_pub.publish(diag)
        self.get_logger().error(f"ESP32-S3 IMU fault: errno={errno}")
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

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

    def destroy_node(self):
        self._write(b"C0,0\n")   # safety: zero on shutdown
        self._close_serial()
        super().destroy_node()


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


if __name__ == "__main__":
    main()
