"""esp32_balance_node.py — ROS2 bridge: Orin ↔ ESP32-S3 BALANCE via USB-CDC serial.

bd-wim1: replaces saltybot_can_bridge (CANable2/python-can) with USB serial.
The CANable2 is freed for Here4 GPS (bd-p47c).

Physical connection
-------------------
  Device  : /dev/esp32-balance  (CH343 USB-CDC, VID:PID 1a86:55d3)
  Baud    : 460800, 8N1
  Protocol: [0xAA][LEN][TYPE][PAYLOAD][CRC8]  (see esp32_balance_protocol.py)

Subscriptions
-------------
  /cmd_vel               geometry_msgs/Twist   → CMD_DRIVE (at up to 50 Hz)
  /estop                 std_msgs/Bool         → CMD_ESTOP
  /saltybot/arm          std_msgs/Bool         → CMD_ARM
  /saltybot/pid_update   std_msgs/String JSON  → CMD_PID  {"kp":…,"ki":…,"kd":…}

Publications
------------
  /saltybot/attitude        std_msgs/String JSON  pitch, motor_cmd, state, etc.
  /saltybot/balance_state   std_msgs/String       same as /saltybot/attitude (alias)
  /can/battery              sensor_msgs/BatteryState  from TELEM_STATUS.vbat_mv
  /can/vesc/left/state      std_msgs/Float32MultiArray  [erpm, duty≡0, voltage_mv/1000, current_ma/1000]
  /can/vesc/right/state     std_msgs/Float32MultiArray  same layout
  /can/connection_status    std_msgs/String       "connected" | "disconnected"

  NOTE: /can/* topic names are preserved for backward compatibility with
        nodes that subscribed to saltybot_can_bridge outputs.

Parameters
----------
  serial_port       str    USB-CDC port           default: /dev/esp32-balance
  baud_rate         int    Baud rate              default: 460800
  speed_scale       float  linear.x (m/s) → units default: 1000.0
  steer_scale       float  angular.z (rad/s) → units  default: -500.0
  command_timeout_s float  /cmd_vel watchdog      default: 0.5
  heartbeat_period  float  CMD_HEARTBEAT interval default: 0.2
  reconnect_delay   float  Retry interval, s      default: 2.0
"""

from __future__ import annotations

import json
import threading
import time
from typing import Optional

import rclpy
import serial
from geometry_msgs.msg import Twist
from rclpy.node import Node
from sensor_msgs.msg import BatteryState
from std_msgs.msg import Bool, Float32MultiArray, String

from .esp32_balance_protocol import (
    BAUD_RATE,
    SERIAL_PORT_DEFAULT,
    AckFrame,
    BalanceState,
    BALANCE_STATE_LABEL,
    FrameParser,
    NackFrame,
    StatusFrame,
    VescFrame,
    encode_arm,
    encode_drive,
    encode_estop,
    encode_heartbeat,
    encode_pid,
)

_WATCHDOG_HZ: float = 20.0   # TX watchdog check rate


class Esp32BalanceNode(Node):
    """Serial bridge replacing CANable2 CAN comms with USB-CDC to ESP32-S3 BALANCE."""

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

        # ── Parameters ────────────────────────────────────────────────────
        self.declare_parameter("serial_port",       SERIAL_PORT_DEFAULT)
        self.declare_parameter("baud_rate",         BAUD_RATE)
        self.declare_parameter("speed_scale",       1000.0)
        self.declare_parameter("steer_scale",       -500.0)
        self.declare_parameter("command_timeout_s", 0.5)
        self.declare_parameter("heartbeat_period",  0.2)
        self.declare_parameter("reconnect_delay",   2.0)

        self._port_name     = self.get_parameter("serial_port").value
        self._baud          = self.get_parameter("baud_rate").value
        self._speed_scale   = self.get_parameter("speed_scale").value
        self._steer_scale   = self.get_parameter("steer_scale").value
        self._cmd_timeout   = self.get_parameter("command_timeout_s").value
        self._hb_period     = self.get_parameter("heartbeat_period").value
        self._reconnect_delay = self.get_parameter("reconnect_delay").value

        # ── State ─────────────────────────────────────────────────────────
        self._ser: Optional[serial.Serial] = None
        self._ser_lock     = threading.Lock()
        self._parser       = FrameParser()
        self._last_cmd_time = time.monotonic()
        self._watchdog_sent = False

        # ── Publishers ────────────────────────────────────────────────────
        self._pub_attitude   = self.create_publisher(String,             "/saltybot/attitude",      10)
        self._pub_balance    = self.create_publisher(String,             "/saltybot/balance_state", 10)
        self._pub_battery    = self.create_publisher(BatteryState,       "/can/battery",            10)
        self._pub_vesc_left  = self.create_publisher(Float32MultiArray,  "/can/vesc/left/state",    10)
        self._pub_vesc_right = self.create_publisher(Float32MultiArray,  "/can/vesc/right/state",   10)
        self._pub_status     = self.create_publisher(String,             "/can/connection_status",  10)

        # ── Subscriptions ─────────────────────────────────────────────────
        self.create_subscription(Twist,  "/cmd_vel",            self._cmd_vel_cb,  10)
        self.create_subscription(Bool,   "/estop",              self._estop_cb,    10)
        self.create_subscription(Bool,   "/saltybot/arm",       self._arm_cb,      10)
        self.create_subscription(String, "/saltybot/pid_update", self._pid_cb,     10)

        # ── Timers ────────────────────────────────────────────────────────
        self.create_timer(self._hb_period,     self._heartbeat_cb)
        self.create_timer(1.0 / _WATCHDOG_HZ,  self._watchdog_cb)
        self.create_timer(self._reconnect_delay, self._reconnect_cb)

        # ── Open port + start reader thread ───────────────────────────────
        self._try_connect()
        self._reader_thread = threading.Thread(
            target=self._reader_loop, daemon=True, name="balance_serial_reader"
        )
        self._reader_thread.start()

        self.get_logger().info(
            f"esp32_balance_node — {self._port_name} @ {self._baud} baud | "
            f"speed_scale={self._speed_scale} steer_scale={self._steer_scale}"
        )

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

    def _try_connect(self) -> None:
        with self._ser_lock:
            if self._ser is not None and self._ser.is_open:
                return
            try:
                self._ser = serial.Serial(
                    port=self._port_name,
                    baudrate=self._baud,
                    timeout=0.05,
                    write_timeout=0.1,
                )
                self._ser.reset_input_buffer()
                self._parser.reset()
                self.get_logger().info(f"Serial open: {self._port_name}")
                self._publish_status("connected")
            except serial.SerialException as exc:
                self.get_logger().warning(
                    f"Cannot open {self._port_name}: {exc} — retry in {self._reconnect_delay:.0f} s",
                    throttle_duration_sec=self._reconnect_delay,
                )
                self._ser = None
                self._publish_status("disconnected")

    def _reconnect_cb(self) -> None:
        with self._ser_lock:
            open_ok = self._ser is not None and self._ser.is_open
        if not open_ok:
            self._try_connect()

    def _handle_serial_error(self, exc: Exception, context: str) -> None:
        self.get_logger().warning(f"Serial error in {context}: {exc}")
        with self._ser_lock:
            if self._ser is not None:
                try:
                    self._ser.close()
                except Exception:
                    pass
            self._ser = None
        self._publish_status("disconnected")

    # ── Serial write helper ───────────────────────────────────────────────

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

    # ── ROS subscriptions ─────────────────────────────────────────────────

    def _cmd_vel_cb(self, msg: Twist) -> None:
        """Convert /cmd_vel to CMD_DRIVE at up to 50 Hz."""
        self._last_cmd_time  = time.monotonic()
        self._watchdog_sent  = False

        speed = int(msg.linear.x  * self._speed_scale)
        steer = int(msg.angular.z * self._steer_scale)
        self._write(encode_drive(speed, steer), "cmd_vel")

    def _estop_cb(self, msg: Bool) -> None:
        frame = encode_estop(msg.data)
        self._write(frame, "estop")
        if msg.data:
            self.get_logger().warning("E-stop asserted → CMD_ESTOP sent to ESP32-S3 BALANCE")

    def _arm_cb(self, msg: Bool) -> None:
        self._write(encode_arm(msg.data), "arm")

    def _pid_cb(self, msg: String) -> None:
        """Parse JSON {"kp":…,"ki":…,"kd":…} and send CMD_PID."""
        try:
            d = json.loads(msg.data)
            kp = float(d["kp"])
            ki = float(d["ki"])
            kd = float(d["kd"])
        except (KeyError, ValueError, json.JSONDecodeError) as exc:
            self.get_logger().error(f"Bad /saltybot/pid_update JSON: {exc}")
            return
        self._write(encode_pid(kp, ki, kd), "pid_update")

    # ── Timers ────────────────────────────────────────────────────────────

    def _heartbeat_cb(self) -> None:
        self._write(encode_heartbeat(), "heartbeat")

    def _watchdog_cb(self) -> None:
        """Send CMD_DRIVE(0,0) if /cmd_vel has been silent for command_timeout_s."""
        if self._watchdog_sent:
            return
        if time.monotonic() - self._last_cmd_time > self._cmd_timeout:
            self._write(encode_drive(0, 0), "watchdog")
            self._watchdog_sent = True

    # ── Background reader ─────────────────────────────────────────────────

    def _reader_loop(self) -> None:
        while rclpy.ok():
            with self._ser_lock:
                ser = self._ser if (self._ser and self._ser.is_open) else None
            if ser is None:
                time.sleep(0.05)
                continue
            try:
                n = ser.in_waiting
                if n:
                    raw = ser.read(n)
                else:
                    time.sleep(0.002)
                    continue
            except serial.SerialException as exc:
                self._handle_serial_error(exc, "reader_loop")
                continue

            for byte in raw:
                frame = self._parser.feed(byte)
                if frame is not None:
                    self._dispatch(frame)

    def _dispatch(self, frame) -> None:
        try:
            if isinstance(frame, StatusFrame):
                self._handle_status(frame)
            elif isinstance(frame, VescFrame):
                self._handle_vesc(frame)
            elif isinstance(frame, AckFrame):
                self.get_logger().debug(
                    f"ACK for cmd type 0x{frame.echoed_type:02X}"
                )
            elif isinstance(frame, NackFrame):
                self.get_logger().warning(
                    f"NACK cmd=0x{frame.cmd_type:02X} err=0x{frame.error_code:02X}"
                )
            elif isinstance(frame, tuple):
                type_code, _ = frame
                self.get_logger().debug(f"Unknown telemetry type 0x{type_code:02X}")
        except Exception as exc:
            self.get_logger().warning(f"Error dispatching frame {type(frame).__name__}: {exc}")

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

    def _handle_status(self, f: StatusFrame) -> None:
        """Publish balance controller status — replaces CAN 0x400 handler."""
        now = self.get_clock().now().to_msg()
        payload = {
            "pitch_deg":    round(f.pitch_deg, 2),
            "motor_cmd":    f.motor_cmd,
            "vbat_mv":      f.vbat_mv,
            "balance_state": f.balance_state,
            "state_label":  BALANCE_STATE_LABEL.get(f.balance_state,
                                                     f"UNKNOWN({f.balance_state})"),
            "flags":        f.flags,
            "estop_active": bool(f.flags & 0x01),
            "hb_timeout":   bool(f.flags & 0x02),
            "ts":           f"{now.sec}.{now.nanosec:09d}",
        }
        msg = String()
        msg.data = json.dumps(payload)
        self._pub_attitude.publish(msg)
        self._pub_balance.publish(msg)

        # Battery piggyback on STATUS frame for nodes that use /can/battery
        bat = BatteryState()
        bat.header.stamp = now
        bat.voltage      = f.vbat_mv / 1000.0
        bat.present      = True
        bat.power_supply_status = BatteryState.POWER_SUPPLY_STATUS_DISCHARGING
        self._pub_battery.publish(bat)

    def _handle_vesc(self, f: VescFrame) -> None:
        """Publish proxied VESC telemetry — replaces CAN VESC_STATUS_1 handler.

        Float32MultiArray layout: [erpm, voltage_v, current_a, temp_c]
        (duty not available via serial proxy; 0.0 placeholder)
        """
        arr = Float32MultiArray()
        arr.data = [
            float(f.erpm),
            float(f.voltage_mv) / 1000.0,
            float(f.current_ma) / 1000.0,
            f.temp_c,
        ]
        if f.vesc_id == 56:
            self._pub_vesc_left.publish(arr)
        elif f.vesc_id == 68:
            self._pub_vesc_right.publish(arr)
        else:
            self.get_logger().warning(f"Unknown VESC ID {f.vesc_id} in telemetry frame")

    # ── Status helper ─────────────────────────────────────────────────────

    def _publish_status(self, status: str) -> None:
        msg = String()
        msg.data = status
        self._pub_status.publish(msg)

    # ── Shutdown ──────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        # Send zero drive + estop before closing
        self._write(encode_drive(0, 0), "shutdown")
        self._write(encode_estop(True), "shutdown")
        with self._ser_lock:
            if self._ser is not None:
                try:
                    self._ser.close()
                except Exception:
                    pass
            self._ser = None
        super().destroy_node()


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


if __name__ == "__main__":
    main()
