"""
mqtt_bridge_node.py — ROS2 ↔ MQTT bridge for cellular telemetry relay.

Forwards key ROS2 topics to a remote MQTT broker over the cellular link,
and receives remote commands from MQTT, republishing them on ROS2.

Bandwidth budget: < 50 KB/s total over cellular.
At 1 Hz all topics: ~1–2 KB/s (well within budget).

ROS2 → MQTT (telemetry, QoS 0):
  /saltybot/imu       → <prefix>/imu         (sensor_msgs/Imu)
  /gps/fix            → <prefix>/gps/fix      (sensor_msgs/NavSatFix)
  /gps/vel            → <prefix>/gps/vel      (geometry_msgs/TwistStamped)
  /uwb/ranges         → <prefix>/uwb/ranges   (std_msgs/String JSON)
  /person/target      → <prefix>/person       (geometry_msgs/PoseStamped)
  /cellular/status    → <prefix>/cellular     (diagnostic_msgs pass-through)

MQTT → ROS2 (commands, QoS 1):
  <prefix>/cmd        → /saltybot/cmd         (std_msgs/String)
  <prefix>/estop      → /saltybot/estop       (std_msgs/Bool)

MQTT topic naming:
  prefix default: "saltybot"
  e.g.  saltybot/gps/fix

Bandwidth throttling:
  Each topic has a configurable max publish rate.
  Default telemetry at 1 Hz, /saltybot/imu at 5 Hz (still manageable).

TLS support:
  Set tls_enable: true and provide tls_ca_cert path for encrypted connection.
"""

import json
import math
import threading
import time

import rclpy
from rclpy.node import Node
from diagnostic_msgs.msg import DiagnosticArray
from geometry_msgs.msg import PoseStamped, TwistStamped
from sensor_msgs.msg import Imu, NavSatFix
from std_msgs.msg import Bool, String

try:
    import paho.mqtt.client as mqtt
    _MQTT_OK = True
except ImportError:
    _MQTT_OK = False


# ── Pure serialisation functions ──────────────────────────────────────────────
# Extracted for unit tests — no ROS2 runtime needed.

def imu_to_dict(msg) -> dict:
    """Serialize sensor_msgs/Imu to bandwidth-minimal dict."""
    return {
        "ax": round(msg.linear_acceleration.x, 3),
        "ay": round(msg.linear_acceleration.y, 3),
        "az": round(msg.linear_acceleration.z, 3),
        "gx": round(msg.angular_velocity.x, 4),
        "gy": round(msg.angular_velocity.y, 4),
        "gz": round(msg.angular_velocity.z, 4),
        "t":  msg.header.stamp.sec,
    }


def navsatfix_to_dict(msg) -> dict:
    """Serialize sensor_msgs/NavSatFix to dict."""
    return {
        "lat":  round(msg.latitude,  7),
        "lon":  round(msg.longitude, 7),
        "alt":  round(msg.altitude,  2),
        "stat": int(msg.status.status),
        "t":    msg.header.stamp.sec,
    }


def gps_vel_to_dict(msg) -> dict:
    """Serialize geometry_msgs/TwistStamped (GPS velocity) to dict."""
    speed = math.sqrt(
        msg.twist.linear.x ** 2 + msg.twist.linear.y ** 2)
    heading_deg = math.degrees(msg.twist.angular.z) % 360.0
    return {
        "spd": round(speed, 2),
        "hdg": round(heading_deg, 1),
        "t":   msg.header.stamp.sec,
    }


def pose_to_dict(msg) -> dict:
    """Serialize geometry_msgs/PoseStamped to dict."""
    return {
        "x":  round(msg.pose.position.x, 3),
        "y":  round(msg.pose.position.y, 3),
        "z":  round(msg.pose.position.z, 3),
        "t":  msg.header.stamp.sec,
    }


# ── Rate limiter helper ───────────────────────────────────────────────────────

class _RateLimiter:
    """True if enough time has elapsed since last allowed call."""
    def __init__(self, rate_hz: float):
        self._interval = 1.0 / rate_hz if rate_hz > 0 else 0.0
        self._last     = 0.0

    def allow(self) -> bool:
        now = time.monotonic()
        if now - self._last >= self._interval:
            self._last = now
            return True
        return False


# ── ROS2 Node ─────────────────────────────────────────────────────────────────

class MqttBridgeNode(Node):

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("mqtt_broker",      "mqtt.saltylab.local")
        self.declare_parameter("mqtt_port",        1883)
        self.declare_parameter("mqtt_user",        "")
        self.declare_parameter("mqtt_password",    "")
        self.declare_parameter("topic_prefix",     "saltybot")
        self.declare_parameter("client_id",        "saltybot-jetson")
        self.declare_parameter("tls_enable",       False)
        self.declare_parameter("tls_ca_cert",      "")
        self.declare_parameter("keepalive",        60)
        self.declare_parameter("reconnect_delay",  5.0)
        # Per-topic rate limits (Hz, 0 = unlimited but still bounded by ROS2 topic rate)
        self.declare_parameter("imu_rate",         5.0)
        self.declare_parameter("gps_rate",         1.0)
        self.declare_parameter("person_rate",      2.0)
        self.declare_parameter("uwb_rate",         1.0)
        self.declare_parameter("cellular_rate",    0.2)

        p = {k: self.get_parameter(k).value for k in [
            "mqtt_broker", "mqtt_port", "mqtt_user", "mqtt_password",
            "topic_prefix", "client_id", "tls_enable", "tls_ca_cert",
            "keepalive", "reconnect_delay",
            "imu_rate", "gps_rate", "person_rate", "uwb_rate", "cellular_rate",
        ]}
        self._p      = p
        self._prefix = p["topic_prefix"].rstrip("/")
        self._client = None
        self._connected = False
        self._lock   = threading.Lock()

        # Rate limiters per outbound topic
        self._rl = {
            "imu":      _RateLimiter(p["imu_rate"]),
            "gps":      _RateLimiter(p["gps_rate"]),
            "person":   _RateLimiter(p["person_rate"]),
            "uwb":      _RateLimiter(p["uwb_rate"]),
            "cellular": _RateLimiter(p["cellular_rate"]),
        }

        if not _MQTT_OK:
            self.get_logger().error(
                "paho-mqtt not installed (pip install paho-mqtt) — bridge inactive")
            return

        # ── MQTT client setup ──────────────────────────────────────────────────
        self._client = mqtt.Client(client_id=p["client_id"])
        if p["mqtt_user"]:
            self._client.username_pw_set(p["mqtt_user"], p["mqtt_password"])
        if p["tls_enable"] and p["tls_ca_cert"]:
            self._client.tls_set(ca_certs=p["tls_ca_cert"])
        self._client.on_connect    = self._on_connect
        self._client.on_disconnect = self._on_disconnect
        self._client.on_message    = self._on_mqtt_message

        self._client.reconnect_delay_set(
            min_delay=1, max_delay=int(p["reconnect_delay"]))

        # Start MQTT loop in background thread
        self._client.connect_async(p["mqtt_broker"], p["mqtt_port"], p["keepalive"])
        self._client.loop_start()

        # ── ROS2 subscriptions (telemetry out) ─────────────────────────────────
        self.create_subscription(Imu,          "/saltybot/imu",    self._imu_cb,    10)
        self.create_subscription(NavSatFix,    "/gps/fix",         self._gps_fix_cb, 10)
        self.create_subscription(TwistStamped, "/gps/vel",         self._gps_vel_cb, 10)
        self.create_subscription(String,       "/uwb/ranges",      self._uwb_cb,    10)
        self.create_subscription(PoseStamped,  "/person/target",   self._person_cb, 10)
        self.create_subscription(DiagnosticArray, "/cellular/status", self._cellular_cb, 5)

        # ── ROS2 publishers (commands in) ──────────────────────────────────────
        self._cmd_pub   = self.create_publisher(String, "/saltybot/cmd",   10)
        self._estop_pub = self.create_publisher(Bool,   "/saltybot/estop", 10)

        self.get_logger().info(
            f"MqttBridge ready  broker={p['mqtt_broker']}:{p['mqtt_port']} "
            f"prefix={self._prefix}")

    # ── MQTT callbacks ────────────────────────────────────────────────────────

    def _on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            self._connected = True
            self.get_logger().info(
                f"MQTT connected to {self._p['mqtt_broker']}")
            # Subscribe to inbound command topics (QoS 1)
            client.subscribe(f"{self._prefix}/cmd",   qos=1)
            client.subscribe(f"{self._prefix}/estop", qos=1)
        else:
            self.get_logger().warn(f"MQTT connect failed rc={rc}")

    def _on_disconnect(self, client, userdata, rc):
        self._connected = False
        if rc != 0:
            self.get_logger().warn(f"MQTT disconnected unexpectedly (rc={rc}), reconnecting")

    def _on_mqtt_message(self, client, userdata, msg):
        """Route inbound MQTT commands to ROS2 topics."""
        try:
            topic   = msg.topic
            payload = msg.payload.decode("utf-8", errors="replace")
            suffix  = topic[len(self._prefix):].lstrip("/")

            if suffix == "cmd":
                ros_msg = String()
                ros_msg.data = payload
                self._cmd_pub.publish(ros_msg)

            elif suffix == "estop":
                try:
                    data = json.loads(payload)
                    val  = bool(data.get("estop", False))
                except (json.JSONDecodeError, AttributeError):
                    val = payload.strip().lower() in ("true", "1", "yes")
                ros_msg = Bool()
                ros_msg.data = val
                self._estop_pub.publish(ros_msg)
                if val:
                    self.get_logger().warn("ESTOP received via MQTT")

        except Exception as exc:
            self.get_logger().error(f"MQTT message error: {exc}")

    # ── Publish helper ────────────────────────────────────────────────────────

    def _mqtt_publish(self, subtopic: str, payload: dict, qos: int = 0):
        """Serialize payload to JSON and publish to MQTT (QoS 0 by default)."""
        if not self._connected or self._client is None:
            return
        try:
            self._client.publish(
                f"{self._prefix}/{subtopic}",
                json.dumps(payload, separators=(",", ":")),
                qos=qos,
            )
        except Exception as exc:
            self.get_logger().warn(f"MQTT publish error: {exc}")

    # ── ROS2 topic callbacks (telemetry out) ──────────────────────────────────

    def _imu_cb(self, msg):
        if self._rl["imu"].allow():
            self._mqtt_publish("imu", imu_to_dict(msg))

    def _gps_fix_cb(self, msg):
        if self._rl["gps"].allow():
            self._mqtt_publish("gps/fix", navsatfix_to_dict(msg))

    def _gps_vel_cb(self, msg):
        if self._rl["gps"].allow():
            self._mqtt_publish("gps/vel", gps_vel_to_dict(msg))

    def _uwb_cb(self, msg):
        if self._rl["uwb"].allow():
            try:
                data = json.loads(msg.data)
            except json.JSONDecodeError:
                data = {"raw": msg.data}
            self._mqtt_publish("uwb/ranges", data)

    def _person_cb(self, msg):
        if self._rl["person"].allow():
            self._mqtt_publish("person", pose_to_dict(msg))

    def _cellular_cb(self, msg):
        if self._rl["cellular"].allow() and msg.status:
            ds = msg.status[0]
            d  = {kv.key: kv.value for kv in ds.values}
            d["level"] = ds.level
            self._mqtt_publish("cellular", d)

    def destroy_node(self):
        if self._client:
            self._client.loop_stop()
            self._client.disconnect()
        super().destroy_node()


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

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


if __name__ == "__main__":
    main()
