"""
saltybot_bridge — serial_bridge_node
<<<<<<< HEAD
ESP32 USB CDC → ROS2 topic publisher
=======
ESP32-S3 USB CDC → ROS2 topic publisher
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

Telemetry frame (50 Hz, newline-delimited JSON):
  {"p":<pitch×10>,"r":<roll×10>,"e":<err×10>,"ig":<integral×10>,
   "m":<motor_cmd -1000..1000>,"s":<state 0-2>,"y":<yaw×10>}

  Error frame (IMU fault):
  {"err":<imu_errno>}

State values (balance_state_t):
  0 = DISARMED  1 = ARMED  2 = TILT_FAULT

Published topics:
  /saltybot/imu          sensor_msgs/Imu         — pitch/roll/yaw as angular velocity
  /saltybot/balance_state  std_msgs/String (JSON)  — full PID diagnostics
"""

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

import serial

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

<<<<<<< HEAD
# Balance state labels matching ESP32 BALANCE balance_state_t enum
=======
# Balance state labels matching ESP32-S3 balance_state_t enum
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
_STATE_LABEL = {0: "DISARMED", 1: "ARMED", 2: "TILT_FAULT"}

# Sensor frame_id published in Imu header
IMU_FRAME_ID = "imu_link"


class SerialBridgeNode(Node):
    def __init__(self):
        super().__init__("esp32_serial_bridge")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("serial_port", "/dev/ttyACM0")
        self.declare_parameter("baud_rate", 921600)
        self.declare_parameter("timeout", 0.1)
        self.declare_parameter("reconnect_delay", 2.0)

        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

        # ── QoS ───────────────────────────────────────────────────────────────
        sensor_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=10,
        )
        diag_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", diag_qos
        )
        self._diag_pub = self.create_publisher(
            DiagnosticArray, "/diagnostics", diag_qos
        )

        # ── State ─────────────────────────────────────────────────────────────
        self._port_name = port
        self._baud = baud
        self._timeout = timeout
        self._ser: serial.Serial | None = None
        self._frame_count = 0
        self._error_count = 0
        self._last_state = -1

        # ── Open serial and start read timer ──────────────────────────────────
        self._open_serial()
<<<<<<< HEAD
        # Poll at 100 Hz — ESP32 BALANCE sends at 50 Hz, so we never miss a frame
=======
        # Poll at 100 Hz — ESP32-S3 sends at 50 Hz, so we never miss a frame
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        self._timer = self.create_timer(0.01, self._read_cb)

        self.get_logger().info(
            f"stm32_serial_bridge started — {port} @ {baud} baud"
        )

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

    def _open_serial(self) -> bool:
        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 port 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):
        if self._ser and self._ser.is_open:
            try:
                self._ser.close()
            except Exception:
                pass
        self._ser = None

    def write_serial(self, data: bytes) -> bool:
        """
<<<<<<< HEAD
        Send raw bytes to ESP32 BALANCE over the open serial port.
=======
        Send raw bytes to ESP32-S3 over the open serial port.
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        Returns False if port is not open (caller should handle gracefully).
        Note: for bidirectional use prefer saltybot_cmd_node which owns TX natively.
        """
        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._close_serial()
            return False

    # ── Read callback ─────────────────────────────────────────────────────────

    def _read_cb(self):
        if self._ser is None or not self._ser.is_open:
            # Attempt reconnect on next timer fire after delay
            self.get_logger().warn(
                f"Serial disconnected — retrying in {self._reconnect_delay}s",
                throttle_duration_sec=self._reconnect_delay,
            )
            self._open_serial()
            return

        try:
            # Read all lines buffered since last call
            while self._ser.in_waiting:
                raw = self._ser.readline()
                if raw:
                    self._parse_and_publish(raw)
        except serial.SerialException as exc:
            self.get_logger().error(f"Serial read error: {exc}")
            self._close_serial()

    # ── Parsing + publishing ──────────────────────────────────────────────────

    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()

        # IMU fault frame — {"err": <errno>}
        if "err" in data:
            self._publish_imu_fault(data["err"], now)
            return

        # Normal telemetry frame — validate required keys
        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

        # Values are ×10 integers from firmware — convert to float
        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"])          # -1000..+1000 ESC units
        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,
        )

        # Log state transitions
        if state != self._last_state:
            label = _STATE_LABEL.get(state, f"UNKNOWN({state})")
            self.get_logger().info(f"Balance state → {label}")
            self._last_state = state

    def _publish_imu(self, pitch_deg, roll_deg, yaw_deg, stamp):
        """
        Publish sensor_msgs/Imu.

<<<<<<< HEAD
        The ESP32 BALANCE IMU gives Euler angles (pitch/roll from accelerometer+gyro
=======
        The ESP32-S3 IMU gives Euler angles (pitch/roll from accelerometer+gyro
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
        fusion, yaw from gyro integration). We publish them as angular_velocity
        for immediate use by slam_toolbox / robot_localization.

        Note: orientation quaternion is left zeroed (covariance [-1,...]) until
        a proper fusion node (e.g. robot_localization EKF) computes it.
        """
        msg = Imu()
        msg.header.stamp = stamp
        msg.header.frame_id = IMU_FRAME_ID

        # orientation unknown — signal with -1 in first covariance element
        msg.orientation_covariance[0] = -1.0

        # angular_velocity: Euler rates approximation (good enough at low freq)
        msg.angular_velocity.x = math.radians(pitch_deg)   # pitch → x
        msg.angular_velocity.y = math.radians(roll_deg)    # roll  → y
        msg.angular_velocity.z = math.radians(yaw_deg)     # yaw   → z

        # angular_velocity covariance (diagonal, degrees converted to rad)
        cov_rad = math.radians(0.5) ** 2  # ±0.5° noise estimate
        msg.angular_velocity_covariance[0] = cov_rad
        msg.angular_velocity_covariance[4] = cov_rad
        msg.angular_velocity_covariance[8] = cov_rad

        # linear_acceleration: unknown from this frame
        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
    ):
        """Publish full PID diagnostics as JSON string and /diagnostics."""
        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),
            "frames":    self._frame_count,
            "parse_errors": self._error_count,
        }

        str_msg = String()
        str_msg.data = json.dumps(payload)
        self._balance_pub.publish(str_msg)

        # /diagnostics
        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:   # ARMED
            status.level = DiagnosticStatus.OK
        elif state == 0:  # DISARMED
            status.level = DiagnosticStatus.WARN
        else:             # TILT_FAULT
            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="state",         value=state_label),
        ]
        diag.status.append(status)
        self._diag_pub.publish(diag)

    def _publish_imu_fault(self, errno: int, stamp):
        """On IMU fault frame, publish an ERROR diagnostic."""
        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 reported 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 reported IMU fault: errno={errno}")
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

    def destroy_node(self):
        self._close_serial()
        super().destroy_node()


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


if __name__ == "__main__":
    main()
