"""
rover_driver_node.py — SaltyRover 4-wheel independent ESC motor driver (Issue #110).

Hardware
────────
  SaltyRover: 4-wheel ground robot with individual brushless ESCs.
  ESCs controlled via PWM (servo-style 1000–2000 µs pulses).
<<<<<<< HEAD
  Communication: USB CDC serial to ESP32 BALANCE or Raspberry Pi Pico GPIO PWM bridge.
=======
  Communication: USB CDC serial to ESP32-S3 or Raspberry Pi Pico GPIO PWM bridge.
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

  ESC channel assignments (configurable):
    CH1 = left-front
    CH2 = left-rear
    CH3 = right-front
    CH4 = right-rear

Drive variants
──────────────
  Selectable at runtime via `four_wheel` parameter (default True).

  2-wheel (four_wheel=False):
    Left motors share a single speed, right motors share a single speed.
    Equivalent to standard differential-drive; only CH1+CH3 are used
    (or all four with mirrored values, depending on hardware).

  4-wheel skid-steer (four_wheel=True):
    All four ESCs commanded independently. Front/rear same-side share speed,
    making this kinematically equivalent to 2-wheel but all four motors driven.

Serial protocol
────────────────
  Command frame (ASCII, newline-terminated):
    P<ch1_us>,<ch2_us>,<ch3_us>,<ch4_us>\n
    Values: 1000–2000 µs (1500 = neutral/brake)

  Heartbeat:
    H\n  every heartbeat_period_s (ESC controller zeros PWM after timeout)

Subscribes
──────────
  /cmd_vel  (geometry_msgs/Twist) — Nav2-compatible velocity command
              linear.x  = forward (m/s)
              angular.z = yaw rate (rad/s, +CCW)

Publishes
─────────
  /saltybot/rover_pwm   (std_msgs/String JSON) — current PWM values (observability)
  /saltybot/rover_odom  (nav_msgs/Odometry)    — wheel-speed odometry

Parameters
──────────
  serial_port        /dev/ttyUSB_rover  (or /dev/ttyACM1, etc.)
  baud_rate          115200
  four_wheel         True   — runtime variant flag (toggle live via ros2 param set)
  track_width        0.28   — left↔right wheel centre distance (m)
  max_speed_ms       2.0    — m/s at full ESC deflection
  max_linear_vel     1.5    — clamp on cmd_vel linear.x (m/s)
  max_angular_vel    3.0    — clamp on cmd_vel angular.z (rad/s)
  pwm_neutral_us     1500   — neutral (brake) PWM (µs)
  pwm_range_us       500    — half-range from neutral (µs) → [1000..2000]
  cmd_vel_timeout_s  0.5    — deadman: zero ESCs if /cmd_vel goes silent
  heartbeat_period_s 0.2
  control_rate       50.0   — Hz
  wheelbase          0.30   — front↔rear axle distance for odometry (m)

Usage
─────
  ros2 launch saltybot_rover_driver rover_driver.launch.py
  ros2 param set /rover_driver four_wheel false   # switch to 2-wheel live
"""

import json
import math
import threading
import time

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

from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
from std_msgs.msg import Header, String

from saltybot_rover_driver.kinematics import compute_wheel_speeds, speed_to_pwm

try:
    import serial
    _SERIAL_OK = True
except ImportError:
    _SERIAL_OK = False


class RoverDriverNode(Node):

    def __init__(self):
        super().__init__("rover_driver")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("serial_port",         "/dev/ttyUSB_rover")
        self.declare_parameter("baud_rate",            115200)
        self.declare_parameter("four_wheel",           True)
        self.declare_parameter("track_width",          0.28)
        self.declare_parameter("max_speed_ms",         2.0)
        self.declare_parameter("max_linear_vel",       1.5)
        self.declare_parameter("max_angular_vel",      3.0)
        self.declare_parameter("pwm_neutral_us",       1500)
        self.declare_parameter("pwm_range_us",         500)
        self.declare_parameter("cmd_vel_timeout_s",    0.5)
        self.declare_parameter("heartbeat_period_s",   0.2)
        self.declare_parameter("control_rate",        50.0)
        self.declare_parameter("wheelbase",            0.30)

        self._p = self._load_params()

        # ── State ─────────────────────────────────────────────────────────────
        self._target_linear_x:  float = 0.0
        self._target_angular_z: float = 0.0
        self._last_cmd_vel_t:   float = 0.0
        self._last_pwm:         tuple = (1500, 1500, 1500, 1500)
        self._odom_x:   float = 0.0
        self._odom_y:   float = 0.0
        self._odom_yaw: float = 0.0
        self._last_odom_t: float = time.monotonic()

        # ── Serial ────────────────────────────────────────────────────────────
        self._ser = None
        self._ser_lock = threading.Lock()
        if _SERIAL_OK:
            self._open_serial()
        else:
            self.get_logger().warn(
                "pyserial not installed — running in simulation mode (no serial I/O)")

        # ── QoS ───────────────────────────────────────────────────────────────
        reliable = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            depth=10,
        )

        # ── Subscriptions ─────────────────────────────────────────────────────
        self.create_subscription(Twist, "/cmd_vel", self._cmd_vel_cb, reliable)

        # ── Publishers ────────────────────────────────────────────────────────
        self._pwm_pub  = self.create_publisher(String,  "/saltybot/rover_pwm",  reliable)
        self._odom_pub = self.create_publisher(Odometry, "/saltybot/rover_odom", reliable)

        # ── Timers ────────────────────────────────────────────────────────────
        rate = self._p["control_rate"]
        self._ctrl_timer = self.create_timer(1.0 / rate,           self._control_cb)
        self._hb_timer   = self.create_timer(self._p["hb_period"], self._heartbeat_cb)

        self.get_logger().info(
            f"RoverDriverNode ready  "
            f"variant={'4WD' if self._p['four_wheel'] else '2WD'}  "
            f"track={self._p['track_width']}m  "
            f"max_speed={self._p['max_speed_ms']}m/s  "
            f"port={self._p['serial_port']}"
        )

    # ── Parameter load ────────────────────────────────────────────────────────

    def _load_params(self) -> dict:
        return {
            "serial_port":    self.get_parameter("serial_port").value,
            "baud_rate":      self.get_parameter("baud_rate").value,
            "four_wheel":     self.get_parameter("four_wheel").value,
            "track_width":    self.get_parameter("track_width").value,
            "max_speed_ms":   self.get_parameter("max_speed_ms").value,
            "max_linear_vel": self.get_parameter("max_linear_vel").value,
            "max_angular_vel":self.get_parameter("max_angular_vel").value,
            "pwm_neutral":    int(self.get_parameter("pwm_neutral_us").value),
            "pwm_range":      int(self.get_parameter("pwm_range_us").value),
            "cmd_timeout":    self.get_parameter("cmd_vel_timeout_s").value,
            "hb_period":      self.get_parameter("heartbeat_period_s").value,
            "control_rate":   self.get_parameter("control_rate").value,
            "wheelbase":      self.get_parameter("wheelbase").value,
        }

    # ── Serial helpers ────────────────────────────────────────────────────────

    def _open_serial(self) -> None:
        with self._ser_lock:
            try:
                self._ser = serial.Serial(
                    self._p["serial_port"], self._p["baud_rate"], timeout=0.05)
                self.get_logger().info(f"Serial open: {self._p['serial_port']}")
            except Exception as exc:
                self.get_logger().error(f"Cannot open serial: {exc}")
                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 Exception as exc:
                self.get_logger().error(f"Serial write: {exc}", throttle_duration_sec=2.0)
                self._ser = None
                return False

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

    def _cmd_vel_cb(self, msg: Twist) -> None:
        p = self._p
        self._target_linear_x  = max(-p["max_linear_vel"],
                                      min(p["max_linear_vel"],  msg.linear.x))
        self._target_angular_z = max(-p["max_angular_vel"],
                                      min(p["max_angular_vel"], msg.angular.z))
        self._last_cmd_vel_t = time.monotonic()

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

    def _control_cb(self) -> None:
        self._p = self._load_params()
        p   = self._p
        now = time.monotonic()

        # Deadman
        cmd_age = (now - self._last_cmd_vel_t) if self._last_cmd_vel_t > 0.0 else 1e9
        if cmd_age > p["cmd_timeout"]:
            lin_x, ang_z = 0.0, 0.0
        else:
            lin_x, ang_z = self._target_linear_x, self._target_angular_z

        # Kinematics → wheel speeds
        speeds = compute_wheel_speeds(
            lin_x, ang_z,
            track_width=p["track_width"],
            four_wheel=p["four_wheel"],
            max_speed_ms=p["max_speed_ms"],
        )

        # Pad to 4 values for protocol; 2-wheel mirrors front to rear
        if len(speeds) == 2:
            vl, vr = speeds
            speeds4 = (vl, vl, vr, vr)
        else:
            speeds4 = speeds   # (vfl, vlr, vrf, vrr)

        # Convert to PWM
        def to_pwm(v):
            return speed_to_pwm(v, p["max_speed_ms"], p["pwm_neutral"], p["pwm_range"])

        pwm = tuple(to_pwm(v) for v in speeds4)
        self._last_pwm = pwm

        # Send P command
        frame = f"P{pwm[0]},{pwm[1]},{pwm[2]},{pwm[3]}\n".encode("ascii")
        self._write(frame)

        # Publish observability
        pwm_msg = String()
        pwm_msg.data = json.dumps({
            "ch1_us": pwm[0], "ch2_us": pwm[1],
            "ch3_us": pwm[2], "ch4_us": pwm[3],
            "variant": "4WD" if p["four_wheel"] else "2WD",
        })
        self._pwm_pub.publish(pwm_msg)

        # Odometry integration (dead reckoning)
        dt = now - self._last_odom_t
        self._last_odom_t = now
        self._integrate_odometry(speeds4[0], speeds4[2], p["track_width"], dt)
        self._publish_odom(now)

    # ── Heartbeat ─────────────────────────────────────────────────────────────

    def _heartbeat_cb(self) -> None:
        self._write(b"H\n")

    # ── Odometry ─────────────────────────────────────────────────────────────

    def _integrate_odometry(
        self, v_left: float, v_right: float, track_width: float, dt: float
    ) -> None:
        """Simple 2D dead-reckoning from left/right motor speeds."""
        lin  = (v_left + v_right) / 2.0
        ang  = (v_right - v_left) / track_width if track_width > 0.0 else 0.0
        self._odom_yaw += ang * dt
        self._odom_x   += lin * math.cos(self._odom_yaw) * dt
        self._odom_y   += lin * math.sin(self._odom_yaw) * dt

    def _publish_odom(self, now: float) -> None:
        yaw = self._odom_yaw
        msg = Odometry()
        msg.header.stamp    = self.get_clock().now().to_msg()
        msg.header.frame_id = "odom"
        msg.child_frame_id  = "base_link"

        msg.pose.pose.position.x = self._odom_x
        msg.pose.pose.position.y = self._odom_y
        msg.pose.pose.orientation.z = math.sin(yaw / 2.0)
        msg.pose.pose.orientation.w = math.cos(yaw / 2.0)

        lin_x = (self._last_pwm[0] + self._last_pwm[2]) / 2.0  # rough proxy
        msg.twist.twist.linear.x  = self._target_linear_x
        msg.twist.twist.angular.z = self._target_angular_z

        self._odom_pub.publish(msg)

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

    def destroy_node(self) -> None:
        self._write(b"P1500,1500,1500,1500\n")   # neutral on shutdown
        super().destroy_node()


# ── Entry point ───────────────────────────────────────────────────────────────

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


if __name__ == "__main__":
    main()
