"""
uwb_driver_node.py — ROS2 node for MaUWB ESP32-S3 DW3000 UWB integration.

Hardware
────────
  • 2× MaUWB ESP32-S3 DW3000 anchors on robot stem (USB → Orin)
    - Anchor-0: port side  (y = +sep/2)
    - Anchor-1: starboard  (y = -sep/2)
  • 1× MaUWB tag per enrolled person (belt clip)

AT command interface (115200 8N1)
──────────────────────────────────
  Query:
    AT+RANGE?\r\n

  Response (from anchors, TWR protocol):
    +RANGE:<anchor_id>,<range_mm>[,<rssi>[,<tag_addr>]]\r\n

  Tag pairing (optional — targets a specific enrolled tag):
    AT+RANGE_ADDR=<tag_addr>\r\n   → anchor only ranges with that tag

Publishes
─────────
  /uwb/ranges   (saltybot_uwb_msgs/UwbRangeArray)  100 Hz — raw anchor ranges
  /uwb/bearing  (saltybot_uwb_msgs/UwbBearing)      10 Hz — Kalman-fused bearing
  /uwb/target   (geometry_msgs/PoseStamped)          10 Hz — triangulated position
                                                              (kept for backwards compat)

Tag pairing
───────────
  Set enrolled_tag_ids to a list of tag address strings (e.g. ["0x1234ABCD"]).
  When non-empty, ranges from unrecognised tags are silently discarded.
  The matched tag address is stamped in UwbRange.tag_id and UwbBearing.tag_id.
  When enrolled_tag_ids is empty, all ranges are accepted (tag_id = "").

Parameters
──────────
  port_a, port_b          serial ports for anchor-0 / anchor-1
  baudrate                115200 (default)
  anchor_separation       centre-to-centre anchor spacing (m)
  anchor_height           anchor mounting height (m)
  tag_height              person tag height (m)
  range_timeout_s         stale-anchor threshold (s)
  max_range_m / min_range_m  validity window (m)
  kf_process_noise        Kalman Q scalar
  kf_meas_noise           Kalman R scalar
  range_rate              Hz — /uwb/ranges publish rate  (default 100)
  bearing_rate            Hz — /uwb/bearing publish rate (default 10)
  enrolled_tag_ids        list[str] — accepted tag addresses; [] = accept all

Usage
─────
  ros2 launch saltybot_uwb uwb.launch.py
  ros2 launch saltybot_uwb uwb.launch.py port_a:=/dev/ttyUSB2 port_b:=/dev/ttyUSB3
"""

import math
import re
import threading
import time

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped
from std_msgs.msg import Header

from saltybot_uwb_msgs.msg import UwbRange, UwbRangeArray, UwbBearing
from saltybot_uwb.ranging_math import triangulate_2anchor, KalmanFilter2D, bearing_from_pos

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


# +RANGE:<id>,<mm>  or  +RANGE:<id>,<mm>,<rssi>  or  +RANGE:<id>,<mm>,<rssi>,<tag>
_RANGE_RE = re.compile(
    r"\+RANGE\s*:\s*(\d+)\s*,\s*(\d+)"
    r"(?:\s*,\s*(-?[\d.]+)"
    r"(?:\s*,\s*([\w:x]+))?"
    r")?",
    re.IGNORECASE,
)


class SerialReader(threading.Thread):
    """
    Background thread: polls one anchor's UART at maximum TWR rate,
    fires callback on every valid +RANGE response.
    Supports optional tag pairing via AT+RANGE_ADDR command.
    """

    def __init__(self, port, baudrate, anchor_id, callback, logger, tag_addr=None):
        super().__init__(daemon=True)
        self._port      = port
        self._baudrate  = baudrate
        self._anchor_id = anchor_id
        self._callback  = callback
        self._logger    = logger
        self._tag_addr  = tag_addr
        self._running   = False
        self._ser       = None

    def run(self):
        self._running = True
        while self._running:
            try:
                self._ser = serial.Serial(
                    self._port, self._baudrate, timeout=0.5
                )
                self._logger.info(
                    f"Anchor-{self._anchor_id}: opened {self._port}"
                    + (f" paired with tag {self._tag_addr}" if self._tag_addr else "")
                )
                if self._tag_addr:
                    self._send_pairing_cmd()
                self._read_loop()
            except Exception as exc:
                self._logger.warn(
                    f"Anchor-{self._anchor_id} serial error: {exc} — retrying in 2s"
                )
                if self._ser and self._ser.is_open:
                    self._ser.close()
                time.sleep(2.0)

    def _send_pairing_cmd(self):
        """Configure the anchor to range only with the paired tag."""
        try:
            cmd = f"AT+RANGE_ADDR={self._tag_addr}\r\n".encode("ascii")
            self._ser.write(cmd)
            time.sleep(0.1)
            self._logger.info(
                f"Anchor-{self._anchor_id}: sent tag pairing {self._tag_addr}"
            )
        except Exception as exc:
            self._logger.warn(
                f"Anchor-{self._anchor_id}: pairing cmd failed: {exc}"
            )

    def _read_loop(self):
        while self._running:
            try:
                self._ser.write(b"AT+RANGE?\r\n")
                for _ in range(10):
                    raw = self._ser.readline()
                    if not raw:
                        break
                    line = raw.decode("ascii", errors="replace").strip()
                    m = _RANGE_RE.search(line)
                    if m:
                        range_mm = int(m.group(2))
                        rssi     = float(m.group(3)) if m.group(3) else 0.0
                        tag_addr = m.group(4) if m.group(4) else ""
                        self._callback(self._anchor_id, range_mm, rssi, tag_addr)
                        break
            except Exception as exc:
                self._logger.warn(
                    f"Anchor-{self._anchor_id} read error: {exc}"
                )
                break

    def stop(self):
        self._running = False
        if self._ser and self._ser.is_open:
            self._ser.close()


class UwbDriverNode(Node):

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

        self.declare_parameter("port_a",             "/dev/uwb-anchor0")
        self.declare_parameter("port_b",             "/dev/uwb-anchor1")
        self.declare_parameter("baudrate",           115200)
        self.declare_parameter("anchor_separation",  0.25)
        self.declare_parameter("anchor_height",      0.80)
        self.declare_parameter("tag_height",         0.90)
        self.declare_parameter("range_timeout_s",    1.0)
        self.declare_parameter("max_range_m",        8.0)
        self.declare_parameter("min_range_m",        0.05)
        self.declare_parameter("kf_process_noise",   0.1)
        self.declare_parameter("kf_meas_noise",      0.3)
        self.declare_parameter("range_rate",         100.0)
        self.declare_parameter("bearing_rate",       10.0)
        self.declare_parameter("enrolled_tag_ids",   [""])

        self._p = self._load_params()

        raw_ids = self.get_parameter("enrolled_tag_ids").value
        self._enrolled_tags = [t.strip() for t in raw_ids if t.strip()]
        paired_tag = self._enrolled_tags[0] if self._enrolled_tags else None

        self._lock   = threading.Lock()
        self._ranges: dict = {}
        self._kf     = KalmanFilter2D(
            process_noise=self._p["kf_process_noise"],
            measurement_noise=self._p["kf_meas_noise"],
            dt=1.0 / self._p["bearing_rate"],
        )

        self._ranges_pub  = self.create_publisher(UwbRangeArray, "/uwb/ranges",  10)
        self._bearing_pub = self.create_publisher(UwbBearing,    "/uwb/bearing", 10)
        self._target_pub  = self.create_publisher(PoseStamped,   "/uwb/target",  10)

        if _SERIAL_AVAILABLE:
            self._reader_a = SerialReader(
                self._p["port_a"], self._p["baudrate"],
                anchor_id=0, callback=self._range_cb,
                logger=self.get_logger(),
                tag_addr=paired_tag,
            )
            self._reader_b = SerialReader(
                self._p["port_b"], self._p["baudrate"],
                anchor_id=1, callback=self._range_cb,
                logger=self.get_logger(),
                tag_addr=paired_tag,
            )
            self._reader_a.start()
            self._reader_b.start()
        else:
            self.get_logger().warn(
                "pyserial not installed — running in simulation mode (no serial I/O)"
            )

        self._range_timer = self.create_timer(
            1.0 / self._p["range_rate"], self._range_publish_cb
        )
        self._bearing_timer = self.create_timer(
            1.0 / self._p["bearing_rate"], self._bearing_publish_cb
        )

        self.get_logger().info(
            f"UWB driver ready  sep={self._p['anchor_separation']}m "
            f"ports={self._p['port_a']},{self._p['port_b']} "
            f"range={self._p['range_rate']}Hz "
            f"bearing={self._p['bearing_rate']}Hz "
            f"enrolled_tags={self._enrolled_tags or ['<any>']}"
        )

    def _load_params(self):
        return {
            "port_a":            self.get_parameter("port_a").value,
            "port_b":            self.get_parameter("port_b").value,
            "baudrate":          self.get_parameter("baudrate").value,
            "anchor_separation": self.get_parameter("anchor_separation").value,
            "anchor_height":     self.get_parameter("anchor_height").value,
            "tag_height":        self.get_parameter("tag_height").value,
            "range_timeout_s":   self.get_parameter("range_timeout_s").value,
            "max_range_m":       self.get_parameter("max_range_m").value,
            "min_range_m":       self.get_parameter("min_range_m").value,
            "kf_process_noise":  self.get_parameter("kf_process_noise").value,
            "kf_meas_noise":     self.get_parameter("kf_meas_noise").value,
            "range_rate":        self.get_parameter("range_rate").value,
            "bearing_rate":      self.get_parameter("bearing_rate").value,
        }

    def _is_enrolled(self, tag_addr: str) -> bool:
        if not self._enrolled_tags:
            return True
        return tag_addr in self._enrolled_tags

    def _range_cb(self, anchor_id: int, range_mm: int, rssi: float, tag_addr: str):
        if not self._is_enrolled(tag_addr):
            return
        range_m = range_mm / 1000.0
        p = self._p
        if range_m < p["min_range_m"] or range_m > p["max_range_m"]:
            return
        with self._lock:
            self._ranges[anchor_id] = (range_m, rssi, tag_addr, time.monotonic())

    def _range_publish_cb(self):
        """100 Hz: publish current raw ranges as UwbRangeArray."""
        now     = time.monotonic()
        timeout = self._p["range_timeout_s"]
        with self._lock:
            valid = {
                aid: entry
                for aid, entry in self._ranges.items()
                if (now - entry[3]) <= timeout
            }
        hdr = Header()
        hdr.stamp    = self.get_clock().now().to_msg()
        hdr.frame_id = "base_link"
        arr = UwbRangeArray()
        arr.header = hdr
        for aid, (r, rssi, tag_id, _) in valid.items():
            entry            = UwbRange()
            entry.header     = hdr
            entry.anchor_id  = aid
            entry.range_m    = float(r)
            entry.raw_mm     = int(round(r * 1000.0))
            entry.rssi       = float(rssi)
            entry.tag_id     = tag_id
            arr.ranges.append(entry)
        self._ranges_pub.publish(arr)

    def _bearing_publish_cb(self):
        """10 Hz: Kalman predict+update, publish fused bearing."""
        now     = time.monotonic()
        timeout = self._p["range_timeout_s"]
        sep     = self._p["anchor_separation"]
        with self._lock:
            valid = {
                aid: entry
                for aid, entry in self._ranges.items()
                if (now - entry[3]) <= timeout
            }
        if not valid:
            return
        both_fresh  = 0 in valid and 1 in valid
        confidence  = 1.0 if both_fresh else 0.5
        active_tag  = valid[next(iter(valid))][2]
        dt = 1.0 / self._p["bearing_rate"]
        self._kf.predict(dt=dt)
        if both_fresh:
            r0 = valid[0][0]
            r1 = valid[1][0]
            try:
                x_t, y_t = triangulate_2anchor(
                    r0=r0, r1=r1, sep=sep,
                    anchor_z=self._p["anchor_height"],
                    tag_z=self._p["tag_height"],
                )
            except (ValueError, ZeroDivisionError) as exc:
                self.get_logger().warn(f"Triangulation error: {exc}")
                return
            self._kf.update(x_t, y_t)
        kx, ky = self._kf.position()
        bearing = bearing_from_pos(kx, ky)
        range_m = math.sqrt(kx * kx + ky * ky)
        hdr = Header()
        hdr.stamp    = self.get_clock().now().to_msg()
        hdr.frame_id = "base_link"
        brg_msg             = UwbBearing()
        brg_msg.header      = hdr
        brg_msg.bearing_rad = float(bearing)
        brg_msg.range_m     = float(range_m)
        brg_msg.confidence  = float(confidence)
        brg_msg.tag_id      = active_tag
        self._bearing_pub.publish(brg_msg)
        pose = PoseStamped()
        pose.header = hdr
        pose.pose.position.x = kx
        pose.pose.position.y = ky
        pose.pose.position.z = 0.0
        yaw = bearing
        pose.pose.orientation.z = math.sin(yaw / 2.0)
        pose.pose.orientation.w = math.cos(yaw / 2.0)
        self._target_pub.publish(pose)


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


if __name__ == "__main__":
    main()
