"""
tank_driver_node.py — SaltyTank tracked-vehicle ESC driver (Issue #122).

Hardware
────────
  SaltyTank: tracked robot with left/right independent 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  (or left-track in 2WD/tracked mode)
    CH2 = left-rear   (mirrored in 2WD/tracked mode)
    CH3 = right-front (or right-track in 2WD/tracked mode)
    CH4 = right-rear  (mirrored in 2WD/tracked mode)

Drive modes
───────────
  Selectable at runtime via `drive_mode` parameter (default "tracked").

  "2wd"     — standard differential drive; CH1+CH2 = left, CH3+CH4 = right.
              No slip compensation.

  "4wd"     — skid-steer; all four ESCs commanded.  Left pair same speed,
              right pair same speed.  No slip compensation.

  "tracked" — tank kinematics with slip compensation.  Angular command is
              pre-scaled to achieve the desired yaw rate despite track slip.
              Odometry angular is de-scaled consistently.

  Switch live: ros2 param set /tank_driver drive_mode tracked

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

  Emergency stop frame:
    E\\n  — forces all channels to neutral immediately

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

Emergency stop
──────────────
  Two stop paths:
    1. Software watchdog  — /cmd_vel silence > watchdog_timeout_s → neutral
    2. E-stop topic       — /saltybot/e_stop (std_msgs/Bool, True = stop)

  Both paths send P1500,1500,1500,1500\\n and latch until a new /cmd_vel arrives
  (e-stop topic) or new commands come in (watchdog).

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

Publishes
─────────
  /saltybot/tank_pwm   (std_msgs/String JSON) — current PWM values
  /saltybot/tank_odom  (nav_msgs/Odometry)    — dead-reckoning odometry

Parameters
──────────
  serial_port        /dev/ttyUSB_tank   (or /dev/ttyACM0, etc.)
  baud_rate          115200
  drive_mode         "tracked"          — "2wd" | "4wd" | "tracked"
  track_width        0.30               — left↔right track centre distance (m)
  wheelbase          0.32               — front↔rear axle distance (m)
  max_speed_ms       1.5                — m/s at full ESC deflection
  max_linear_vel     1.0                — clamp on cmd_vel linear.x (m/s)
  max_angular_vel    2.5                — clamp on cmd_vel angular.z (rad/s)
  slip_factor        0.3                — track slip coefficient [0.0, 0.99]
  pwm_neutral_us     1500
  pwm_range_us       500
  watchdog_timeout_s 0.3                — tighter than rover (0.5 s) for safety
  heartbeat_period_s 0.2
  control_rate       50.0               — Hz

Usage
─────
  ros2 launch saltybot_tank_driver tank_driver.launch.py
  ros2 param set /tank_driver drive_mode 4wd
  ros2 param set /tank_driver slip_factor 0.1
  ros2 topic pub /saltybot/e_stop std_msgs/Bool '{data: true}' --once
"""

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 Bool, String

from saltybot_tank_driver.kinematics import (
    compute_track_speeds,
    expand_to_4ch,
    odometry_from_track_speeds,
    speed_to_pwm,
    DRIVE_MODES,
)

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


class TankDriverNode(Node):

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

        # ── Parameters ───────────────────────────────────────────────────────
        self.declare_parameter("serial_port",        "/dev/ttyUSB_tank")
        self.declare_parameter("baud_rate",          115200)
        self.declare_parameter("drive_mode",         "tracked")
        self.declare_parameter("track_width",        0.30)
        self.declare_parameter("wheelbase",          0.32)
        self.declare_parameter("max_speed_ms",       1.5)
        self.declare_parameter("max_linear_vel",     1.0)
        self.declare_parameter("max_angular_vel",    2.5)
        self.declare_parameter("slip_factor",        0.3)
        self.declare_parameter("pwm_neutral_us",     1500)
        self.declare_parameter("pwm_range_us",       500)
        self.declare_parameter("watchdog_timeout_s", 0.3)
        self.declare_parameter("heartbeat_period_s", 0.2)
        self.declare_parameter("control_rate",       50.0)

        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._e_stop:           bool  = False

        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)
        self.create_subscription(Bool,  "/saltybot/e_stop", self._e_stop_cb,  reliable)

        # ── Publishers ───────────────────────────────────────────────────────
        self._pwm_pub  = self.create_publisher(String,   "/saltybot/tank_pwm",  reliable)
        self._odom_pub = self.create_publisher(Odometry, "/saltybot/tank_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"TankDriverNode ready  "
            f"mode={self._p['drive_mode']}  "
            f"track={self._p['track_width']}m  "
            f"slip={self._p['slip_factor']}  "
            f"max_speed={self._p['max_speed_ms']}m/s  "
            f"port={self._p['serial_port']}"
        )

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

    def _load_params(self) -> dict:
        mode = self.get_parameter("drive_mode").value
        if mode not in DRIVE_MODES:
            self.get_logger().warn(
                f"Unknown drive_mode '{mode}', defaulting to 'tracked'")
            mode = "tracked"
        return {
            "serial_port":    self.get_parameter("serial_port").value,
            "baud_rate":      self.get_parameter("baud_rate").value,
            "drive_mode":     mode,
            "track_width":    self.get_parameter("track_width").value,
            "wheelbase":      self.get_parameter("wheelbase").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,
            "slip_factor":    float(self.get_parameter("slip_factor").value),
            "pwm_neutral":    int(self.get_parameter("pwm_neutral_us").value),
            "pwm_range":      int(self.get_parameter("pwm_range_us").value),
            "watchdog_timeout": self.get_parameter("watchdog_timeout_s").value,
            "hb_period":      self.get_parameter("heartbeat_period_s").value,
            "control_rate":   self.get_parameter("control_rate").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

    def _send_neutral(self) -> None:
        self._write(b"P1500,1500,1500,1500\n")
        self._last_pwm = (1500, 1500, 1500, 1500)

    # ── 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()
        # A new cmd_vel clears a latched e-stop
        if self._e_stop:
            self._e_stop = False
            self.get_logger().info("E-stop cleared by new cmd_vel")

    # ── E-stop callback ───────────────────────────────────────────────────────

    def _e_stop_cb(self, msg: Bool) -> None:
        if msg.data and not self._e_stop:
            self._e_stop = True
            self._send_neutral()
            self.get_logger().warn("Emergency stop engaged")
        elif not msg.data and self._e_stop:
            self._e_stop = False
            self.get_logger().info("E-stop released via topic")

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

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

        # E-stop or watchdog
        cmd_age = (now - self._last_cmd_vel_t) if self._last_cmd_vel_t > 0.0 else 1e9
        if self._e_stop or cmd_age > p["watchdog_timeout"]:
            lin_x, ang_z = 0.0, 0.0
        else:
            lin_x, ang_z = self._target_linear_x, self._target_angular_z

        # Kinematics → left/right track speeds
        vl, vr = compute_track_speeds(
            lin_x, ang_z,
            track_width=p["track_width"],
            max_speed_ms=p["max_speed_ms"],
            drive_mode=p["drive_mode"],
            slip_factor=p["slip_factor"],
        )

        # Expand to 4 channels
        speeds4 = expand_to_4ch(vl, vr)   # (vl, vl, vr, vr)

        # 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],
            "drive_mode": p["drive_mode"],
            "e_stop": self._e_stop,
        })
        self._pwm_pub.publish(pwm_msg)

        # Odometry integration (dead reckoning with slip compensation)
        dt = now - self._last_odom_t
        self._last_odom_t = now
        self._integrate_odometry(vl, vr, p["track_width"], p["slip_factor"], dt)
        self._publish_odom()

    # ── 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,
        slip_factor: float,
        dt: float,
    ) -> None:
        """2D dead-reckoning with track-slip compensation."""
        lin, ang = odometry_from_track_speeds(v_left, v_right, track_width, slip_factor)
        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) -> 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)

        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._send_neutral()      # safe-state on shutdown
        super().destroy_node()


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

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


if __name__ == "__main__":
    main()
