"""stm32_cmd_node.py — Full bidirectional binary-framed STM32↔Jetson bridge.

Issue #119: replaces the ASCII-protocol saltybot_cmd_node with a robust binary
framing protocol (STX/TYPE/LEN/PAYLOAD/CRC16/ETX) at 921600 baud.

TX commands (Jetson → STM32):
  SPEED_STEER  — 50 Hz from /cmd_vel subscription
  HEARTBEAT    — 200 ms timer (STM32 watchdog fires at 500 ms)
  ARM          — via /saltybot/arm service
  SET_MODE     — via /saltybot/set_mode service
  PID_UPDATE   — via /saltybot/pid_update topic

Watchdog: if /cmd_vel is silent for 500 ms, send SPEED_STEER(0,0) and log warning.

RX telemetry (STM32 → Jetson):
  IMU          → /saltybot/imu           (sensor_msgs/Imu)
  BATTERY      → /saltybot/telemetry/battery (std_msgs/String JSON)
  MOTOR_RPM    → /saltybot/telemetry/motor_rpm (std_msgs/String JSON)
  ARM_STATE    → /saltybot/arm_state     (std_msgs/String JSON)
  ERROR        → /saltybot/error         (std_msgs/String JSON)
  All frames   → /diagnostics            (diagnostic_msgs/DiagnosticArray)

Auto-reconnect: USB disconnect is detected when serial.read() raises; node
continuously retries at reconnect_delay interval.

This node owns /dev/ttyACM0 exclusively — do NOT run alongside
serial_bridge_node or saltybot_cmd_node on the same port.

Parameters (config/stm32_cmd_params.yaml):
  serial_port       /dev/ttyACM0
  baud_rate         921600
  reconnect_delay   2.0   (seconds)
  heartbeat_period  0.2   (seconds)
  watchdog_timeout  0.5   (seconds — no /cmd_vel → send zero-speed)
  speed_scale       1000.0 (linear.x m/s → ESC units)
  steer_scale      -500.0  (angular.z rad/s → ESC units, neg to flip convention)
"""

from __future__ import annotations

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
from std_srvs.srv import SetBool, Trigger

from .stm32_protocol import (
    FrameParser,
    ImuFrame, BatteryFrame, MotorRpmFrame, ArmStateFrame, ErrorFrame,
    encode_heartbeat, encode_speed_steer, encode_arm, encode_set_mode,
    encode_pid_update,
)

# ── Constants ─────────────────────────────────────────────────────────────────

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


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


# ── Node ──────────────────────────────────────────────────────────────────────

class Stm32CmdNode(Node):
    """Binary-framed Jetson↔STM32 bridge node."""

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("serial_port",      "/dev/ttyACM0")
        self.declare_parameter("baud_rate",        921600)
        self.declare_parameter("reconnect_delay",  2.0)
        self.declare_parameter("heartbeat_period", 0.2)
        self.declare_parameter("watchdog_timeout", 0.5)
        self.declare_parameter("speed_scale",      1000.0)
        self.declare_parameter("steer_scale",      -500.0)

        port                  = self.get_parameter("serial_port").value
        baud                  = self.get_parameter("baud_rate").value
        self._reconnect_delay = self.get_parameter("reconnect_delay").value
        self._hb_period       = self.get_parameter("heartbeat_period").value
        self._wd_timeout      = self.get_parameter("watchdog_timeout").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,
        )
        rel_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST, depth=10,
        )

        # ── Publishers ────────────────────────────────────────────────────────
        self._imu_pub      = self.create_publisher(Imu,    "/saltybot/imu",           sensor_qos)
        self._arm_pub      = self.create_publisher(String, "/saltybot/arm_state",      rel_qos)
        self._error_pub    = self.create_publisher(String, "/saltybot/error",          rel_qos)
        self._battery_pub  = self.create_publisher(String, "/saltybot/telemetry/battery",  rel_qos)
        self._rpm_pub      = self.create_publisher(String, "/saltybot/telemetry/motor_rpm", rel_qos)
        self._diag_pub     = self.create_publisher(DiagnosticArray, "/diagnostics",    rel_qos)

        # ── Subscribers ───────────────────────────────────────────────────────
        self._cmd_vel_sub = self.create_subscription(
            Twist, "/cmd_vel", self._on_cmd_vel, rel_qos,
        )
        self._pid_sub = self.create_subscription(
            String, "/saltybot/pid_update", self._on_pid_update, rel_qos,
        )

        # ── Services ──────────────────────────────────────────────────────────
        self._arm_srv  = self.create_service(SetBool, "/saltybot/arm",      self._svc_arm)
        self._mode_srv = self.create_service(SetBool, "/saltybot/set_mode", self._svc_set_mode)

        # ── Serial state ──────────────────────────────────────────────────────
        self._port_name = port
        self._baud      = baud
        self._ser: serial.Serial | None = None
        self._ser_lock  = threading.Lock()
        self._parser    = FrameParser()

        # ── TX state ──────────────────────────────────────────────────────────
        self._last_speed    = 0
        self._last_steer    = 0
        self._last_cmd_t    = time.monotonic()
        self._watchdog_sent = False     # tracks whether we already sent zero

        # ── Diagnostics state ──────────────────────────────────────────────────
        self._last_arm_state   = -1
        self._last_battery_mv  = 0
        self._rx_frame_count   = 0

        # ── Open serial and start timers ──────────────────────────────────────
        self._open_serial()

        # Read at 200 Hz (serial RX thread is better, but timer keeps ROS2 integration clean)
        self._read_timer = self.create_timer(0.005, self._read_cb)
        # Heartbeat TX
        self._hb_timer   = self.create_timer(self._hb_period, self._heartbeat_cb)
        # Watchdog check (fires at 2× watchdog_timeout for quick detection)
        self._wd_timer   = self.create_timer(self._wd_timeout / 2, self._watchdog_cb)
        # Periodic diagnostics
        self._diag_timer = self.create_timer(1.0, self._publish_diagnostics)

        self.get_logger().info(
            f"stm32_cmd_node started — {port} @ {baud} baud | "
            f"HB {int(self._hb_period * 1000)}ms | WD {int(self._wd_timeout * 1000)}ms"
        )

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

    def _open_serial(self) -> bool:
        with self._ser_lock:
            try:
                self._ser = serial.Serial(
                    port=self._port_name,
                    baudrate=self._baud,
                    timeout=0.005,          # non-blocking reads
                    write_timeout=0.1,
                )
                self._ser.reset_input_buffer()
                self._parser.reset()
                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}",
                    throttle_duration_sec=self._reconnect_delay,
                )
                self._ser = None
                return False

    def _close_serial(self) -> None:
        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 is 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 — read callback ────────────────────────────────────────────────────

    def _read_cb(self) -> None:
        """Read bytes from serial and feed them to the frame parser."""
        raw: bytes | None = None
        reconnect_needed  = False

        with self._ser_lock:
            if self._ser is None or not self._ser.is_open:
                reconnect_needed = True
            else:
                try:
                    n = self._ser.in_waiting
                    if n:
                        raw = self._ser.read(n)
                except serial.SerialException as exc:
                    self.get_logger().error(f"Serial read error: {exc}")
                    self._ser = None
                    reconnect_needed = True

        if reconnect_needed:
            self.get_logger().warn(
                "Serial disconnected — will retry",
                throttle_duration_sec=self._reconnect_delay,
            )
            time.sleep(self._reconnect_delay)
            self._open_serial()
            return

        if not raw:
            return

        for byte in raw:
            frame = self._parser.feed(byte)
            if frame is not None:
                self._rx_frame_count += 1
                self._dispatch_frame(frame)

    def _dispatch_frame(self, frame) -> None:
        """Route a decoded frame to the appropriate publisher."""
        now = self.get_clock().now().to_msg()

        if isinstance(frame, ImuFrame):
            self._publish_imu(frame, now)

        elif isinstance(frame, BatteryFrame):
            self._publish_battery(frame, now)

        elif isinstance(frame, MotorRpmFrame):
            self._publish_motor_rpm(frame, now)

        elif isinstance(frame, ArmStateFrame):
            self._publish_arm_state(frame, now)

        elif isinstance(frame, ErrorFrame):
            self._publish_error(frame, now)

        elif isinstance(frame, tuple):
            type_code, payload = frame
            self.get_logger().debug(
                f"Unknown telemetry type 0x{type_code:02X} len={len(payload)}"
            )

    # ── Telemetry publishers ──────────────────────────────────────────────────

    def _publish_imu(self, frame: ImuFrame, stamp) -> None:
        msg = Imu()
        msg.header.stamp    = stamp
        msg.header.frame_id = IMU_FRAME_ID

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

        msg.angular_velocity.x = math.radians(frame.pitch_deg)
        msg.angular_velocity.y = math.radians(frame.roll_deg)
        msg.angular_velocity.z = math.radians(frame.yaw_deg)
        cov = math.radians(0.3) ** 2    # ±0.3° noise estimate from STM32 BMI088
        msg.angular_velocity_covariance[0] = cov
        msg.angular_velocity_covariance[4] = cov
        msg.angular_velocity_covariance[8] = cov

        msg.linear_acceleration.x = frame.accel_x
        msg.linear_acceleration.y = frame.accel_y
        msg.linear_acceleration.z = frame.accel_z
        acov = 0.05 ** 2    # ±0.05 m/s² noise
        msg.linear_acceleration_covariance[0] = acov
        msg.linear_acceleration_covariance[4] = acov
        msg.linear_acceleration_covariance[8] = acov

        self._imu_pub.publish(msg)

    def _publish_battery(self, frame: BatteryFrame, stamp) -> None:
        payload = {
            "voltage_v":  round(frame.voltage_mv / 1000.0, 3),
            "voltage_mv": frame.voltage_mv,
            "current_ma": frame.current_ma,
            "soc_pct":    frame.soc_pct,
            "charging":   frame.current_ma < -100,
            "ts":         f"{stamp.sec}.{stamp.nanosec:09d}",
        }
        self._last_battery_mv = frame.voltage_mv
        msg = String()
        msg.data = json.dumps(payload)
        self._battery_pub.publish(msg)

    def _publish_motor_rpm(self, frame: MotorRpmFrame, stamp) -> None:
        payload = {
            "left_rpm":  frame.left_rpm,
            "right_rpm": frame.right_rpm,
            "ts":        f"{stamp.sec}.{stamp.nanosec:09d}",
        }
        msg = String()
        msg.data = json.dumps(payload)
        self._rpm_pub.publish(msg)

    def _publish_arm_state(self, frame: ArmStateFrame, stamp) -> None:
        label = _ARM_LABEL.get(frame.state, f"UNKNOWN({frame.state})")
        if frame.state != self._last_arm_state:
            self.get_logger().info(f"Arm state → {label} (flags=0x{frame.error_flags:02X})")
            self._last_arm_state = frame.state

        payload = {
            "state":       frame.state,
            "state_label": label,
            "error_flags": frame.error_flags,
            "ts":          f"{stamp.sec}.{stamp.nanosec:09d}",
        }
        msg = String()
        msg.data = json.dumps(payload)
        self._arm_pub.publish(msg)

    def _publish_error(self, frame: ErrorFrame, stamp) -> None:
        self.get_logger().error(
            f"STM32 error code=0x{frame.error_code:02X} sub=0x{frame.subcode:02X}"
        )
        payload = {
            "error_code": frame.error_code,
            "subcode":    frame.subcode,
            "ts":         f"{stamp.sec}.{stamp.nanosec:09d}",
        }
        msg = String()
        msg.data = json.dumps(payload)
        self._error_pub.publish(msg)

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

    def _on_cmd_vel(self, msg: Twist) -> None:
        """Convert /cmd_vel Twist to SPEED_STEER frame at up to 50 Hz."""
        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
        self._last_cmd_t    = time.monotonic()
        self._watchdog_sent = False

        frame = encode_speed_steer(speed, steer)
        if not self._write(frame):
            self.get_logger().warn(
                "SPEED_STEER dropped — serial not open",
                throttle_duration_sec=2.0,
            )

    def _heartbeat_cb(self) -> None:
        """Send HEARTBEAT every heartbeat_period (default 200ms)."""
        self._write(encode_heartbeat())

    def _watchdog_cb(self) -> None:
        """Send zero-speed if /cmd_vel silent for watchdog_timeout seconds."""
        if time.monotonic() - self._last_cmd_t >= self._wd_timeout:
            if not self._watchdog_sent:
                self.get_logger().warn(
                    f"No /cmd_vel for {self._wd_timeout:.1f}s — sending zero-speed"
                )
                self._watchdog_sent = True
            self._last_speed = 0
            self._last_steer = 0
            self._write(encode_speed_steer(0, 0))

    def _on_pid_update(self, msg: String) -> None:
        """Parse JSON /saltybot/pid_update and send PID_UPDATE frame."""
        try:
            data = json.loads(msg.data)
            kp = float(data["kp"])
            ki = float(data["ki"])
            kd = float(data["kd"])
        except (ValueError, KeyError, json.JSONDecodeError) as exc:
            self.get_logger().error(f"Bad PID update JSON: {exc}")
            return
        frame = encode_pid_update(kp, ki, kd)
        if self._write(frame):
            self.get_logger().info(f"PID update: kp={kp}, ki={ki}, kd={kd}")
        else:
            self.get_logger().warn("PID_UPDATE dropped — serial not open")

    # ── Services ──────────────────────────────────────────────────────────────

    def _svc_arm(self, request: SetBool.Request, response: SetBool.Response):
        """SetBool(True) = arm, SetBool(False) = disarm."""
        arm = request.data
        frame = encode_arm(arm)
        ok = self._write(frame)
        response.success = ok
        response.message = ("ARMED" if arm else "DISARMED") if ok else "serial not open"
        self.get_logger().info(
            f"ARM service: {'arm' if arm else 'disarm'} — {'sent' if ok else 'FAILED'}"
        )
        return response

    def _svc_set_mode(self, request: SetBool.Request, response: SetBool.Response):
        """SetBool: data maps to mode byte (True=1, False=0)."""
        mode = 1 if request.data else 0
        frame = encode_set_mode(mode)
        ok = self._write(frame)
        response.success = ok
        response.message = f"mode={mode}" if ok else "serial not open"
        return response

    # ── Diagnostics ───────────────────────────────────────────────────────────

    def _publish_diagnostics(self) -> None:
        diag = DiagnosticArray()
        diag.header.stamp = self.get_clock().now().to_msg()

        status = DiagnosticStatus()
        status.name        = "saltybot/stm32_cmd_node"
        status.hardware_id = "stm32f722"

        port_ok = self._ser is not None and self._ser.is_open
        if port_ok:
            status.level   = DiagnosticStatus.OK
            status.message = "Serial OK"
        else:
            status.level   = DiagnosticStatus.ERROR
            status.message = f"Serial disconnected: {self._port_name}"

        wd_age = time.monotonic() - self._last_cmd_t
        status.values = [
            KeyValue(key="serial_port",    value=self._port_name),
            KeyValue(key="port_open",      value=str(port_ok)),
            KeyValue(key="rx_frames",      value=str(self._rx_frame_count)),
            KeyValue(key="rx_errors",      value=str(self._parser.frames_error)),
            KeyValue(key="last_speed",     value=str(self._last_speed)),
            KeyValue(key="last_steer",     value=str(self._last_steer)),
            KeyValue(key="cmd_vel_age_s",  value=f"{wd_age:.2f}"),
            KeyValue(key="battery_mv",     value=str(self._last_battery_mv)),
            KeyValue(key="arm_state",      value=_ARM_LABEL.get(self._last_arm_state, "?")),
        ]
        diag.status.append(status)
        self._diag_pub.publish(diag)

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

    def destroy_node(self) -> None:
        # Send zero-speed + disarm on shutdown
        self._write(encode_speed_steer(0, 0))
        self._write(encode_arm(False))
        self._close_serial()
        super().destroy_node()


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


if __name__ == "__main__":
    main()
