"""
<<<<<<< HEAD
saltybot_cmd_node — full bidirectional ESP32 BALANCE↔Jetson bridge
=======
saltybot_cmd_node — full bidirectional ESP32-S3↔Jetson bridge
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
Combines telemetry RX (from serial_bridge_node) with drive command TX.

Owns /dev/ttyACM0 exclusively — do NOT run alongside serial_bridge_node.

<<<<<<< HEAD
RX path (50Hz from ESP32 BALANCE):
  JSON telemetry → /saltybot/imu, /saltybot/balance_state, /diagnostics

TX path:
  /cmd_vel (geometry_msgs/Twist) → C<speed>,<steer>\\n  → ESP32 BALANCE
  Heartbeat timer (200ms)        → H\\n                 → ESP32 BALANCE

Protocol:
  H\\n              — heartbeat. ESP32 BALANCE reverts steer to 0 if gap > 500ms.
  C<spd>,<str>\\n   — drive command. speed/steer: -1000..+1000 integers.
                      C command also refreshes ESP32 BALANCE heartbeat timer.
=======
RX path (50Hz from ESP32-S3):
  JSON telemetry → /saltybot/imu, /saltybot/balance_state, /diagnostics

TX path:
  /cmd_vel (geometry_msgs/Twist) → C<speed>,<steer>\\n  → ESP32-S3
  Heartbeat timer (200ms)        → H\\n                 → ESP32-S3

Protocol:
  H\\n              — heartbeat. ESP32-S3 reverts steer to 0 if gap > 500ms.
  C<spd>,<str>\\n   — drive command. speed/steer: -1000..+1000 integers.
                      C command also refreshes ESP32-S3 heartbeat timer.
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

Twist mapping (configurable via ROS2 params):
  speed = clamp(linear.x  * speed_scale, -1000, 1000)
  steer = clamp(angular.z * steer_scale, -1000, 1000)
  Default scales: speed_scale=1000.0 (1 m/s → 1000), steer_scale=-500.0
  Negative steer_scale because ROS2 +angular.z = CCW but ESC steer convention
  may differ — tune in config/bridge_params.yaml.
"""

import json
import math
import threading
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

import serial

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

_STATE_LABEL = {0: "DISARMED", 1: "ARMED", 2: "TILT_FAULT"}
IMU_FRAME_ID = "imu_link"


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


class SaltybotCmdNode(Node):
    def __init__(self):
        super().__init__("saltybot_cmd")

        # ── 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)   # seconds
        self.declare_parameter("speed_scale",      1000.0) # m/s → ESC units
        self.declare_parameter("steer_scale",      -500.0) # rad/s → ESC units

        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._speed_scale      = self.get_parameter("speed_scale").value
        self._steer_scale      = self.get_parameter("steer_scale").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 (telemetry RX path) ────────────────────────────────────
        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)

        # ── Subscriber (cmd TX path) ──────────────────────────────────────────
        self._cmd_vel_sub = self.create_subscription(
            Twist, "/cmd_vel", self._cmd_vel_cb, reliable_qos)

        # ── State ─────────────────────────────────────────────────────────────
        self._port_name  = port
        self._baud       = baud
        self._timeout    = timeout
        self._ser: serial.Serial | None = None
        self._ser_lock   = threading.Lock()   # guards _ser for RX/TX threads
        self._frame_count  = 0
        self._error_count  = 0
        self._last_state   = -1
        self._last_speed   = 0
        self._last_steer   = 0

        # ── Open serial ───────────────────────────────────────────────────────
        self._open_serial()

        # ── Timers ────────────────────────────────────────────────────────────
<<<<<<< HEAD
        # Telemetry read at 100Hz (ESP32 BALANCE sends at 50Hz)
=======
        # Telemetry read at 100Hz (ESP32-S3 sends at 50Hz)
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        self._read_timer = self.create_timer(0.01, self._read_cb)
        # Heartbeat TX at configured period (default 200ms)
        self._hb_timer   = self.create_timer(self._hb_period, self._heartbeat_cb)

        self.get_logger().info(
            f"saltybot_cmd started — {port} @ {baud} baud "
            f"(HB {int(self._hb_period*1000)}ms, "
            f"speed_scale={self._speed_scale}, steer_scale={self._steer_scale})"
        )

    # ── 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:
        """Thread-safe serial write. Returns False if port not open."""
        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

    # ── RX — telemetry read ───────────────────────────────────────────────────

    def _read_cb(self):
        lines = []
        with self._ser_lock:
            if self._ser is None or not self._ser.is_open:
                pass
            else:
                try:
                    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 = []

        # Parse outside the lock
        for raw in lines:
            self._parse_and_publish(raw)

        # Reconnect if port lost
        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"])
        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, now,
        )
        if state != self._last_state:
            self.get_logger().info(
                f"Balance state → {_STATE_LABEL.get(state, f'UNKNOWN({state})')}"
            )
            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, stamp
    ):
        state_label = _STATE_LABEL.get(state, f"UNKNOWN({state})")
        payload = {
            "stamp":        f"{stamp.sec}.{stamp.nanosec:09d}",
            "state":        state_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),
            "jetson_speed": self._last_speed,
            "jetson_steer": self._last_steer,
            "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 = state_label
        if state == 1:
            status.level = DiagnosticStatus.OK
        elif state == 0:
            status.level = DiagnosticStatus.WARN
        else:
            status.level = 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_deg",  value=f"{error:.1f}"),
            KeyValue(key="integral",       value=f"{integral:.1f}"),
            KeyValue(key="motor_cmd",      value=f"{int(motor_cmd)}"),
            KeyValue(key="jetson_speed",   value=str(self._last_speed)),
            KeyValue(key="jetson_steer",   value=str(self._last_steer)),
            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)

    # ── TX — command send ─────────────────────────────────────────────────────

    def _cmd_vel_cb(self, msg: Twist):
        """Convert Twist to C<speed>,<steer>\\n and send over serial."""
        speed = int(_clamp(msg.linear.x  * self._speed_scale, -1000.0, 1000.0))
        steer = int(_clamp(msg.angular.z * self._steer_scale, -1000.0, 1000.0))
        self._last_speed = speed
        self._last_steer = steer
        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,
            )

    def _heartbeat_cb(self):
<<<<<<< HEAD
        """Send H\\n heartbeat. ESP32 BALANCE reverts steer to 0 if gap > 500ms."""
=======
        """Send H\\n heartbeat. ESP32-S3 reverts steer to 0 if gap > 500ms."""
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        self._write(b"H\n")

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

    def destroy_node(self):
        # Send zero command on shutdown so robot doesn't run away
        self._write(b"C0,0\n")
        self._close_serial()
        super().destroy_node()


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


if __name__ == "__main__":
    main()
