#!/usr/bin/env python3
"""
vesc_mqtt_relay_node.py — ROS2 → MQTT relay for VESC CAN telemetry (Issue #656)

Subscribes to VESC telemetry ROS2 topics and republishes as MQTT JSON payloads
for the sensor dashboard.  Each per-motor topic is rate-limited to 5 Hz.

ROS2 topics consumed
─────────────────────
  /vesc/left/state   std_msgs/String  JSON from vesc_telemetry_node
  /vesc/right/state  std_msgs/String  JSON from vesc_telemetry_node
  /vesc/combined     std_msgs/String  JSON from vesc_telemetry_node

MQTT topics published
──────────────────────
  saltybot/phone/vesc_left      — per-motor telemetry (left)
  saltybot/phone/vesc_right     — per-motor telemetry (right)
  saltybot/phone/vesc_combined  — combined voltage + total current + RPMs

MQTT payload (per-motor)
────────────────────────
  {
    "rpm":          int,
    "current_a":    float,   # phase current
    "voltage_v":    float,   # bus voltage
    "temperature_c": float,  # FET temperature
    "duty_cycle":   float,   # -1.0 … 1.0
    "fault_code":   int,
    "ts":           float    # epoch seconds
  }

Parameters
──────────
  mqtt_host          str   Broker IP/hostname         (default: localhost)
  mqtt_port          int   Broker port                (default: 1883)
  mqtt_keepalive     int   Keepalive seconds          (default: 60)
  reconnect_delay_s  float Delay between reconnects   (default: 5.0)
  motor_rate_hz      float Max publish rate per motor (default: 5.0)

Dependencies
────────────
  pip install paho-mqtt
"""

import json
import time
from typing import Any

import rclpy
from rclpy.node import Node
from std_msgs.msg import String

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

# ── MQTT topic constants ──────────────────────────────────────────────────────

_MQTT_VESC_LEFT     = "saltybot/phone/vesc_left"
_MQTT_VESC_RIGHT    = "saltybot/phone/vesc_right"
_MQTT_VESC_COMBINED = "saltybot/phone/vesc_combined"

# ── ROS2 topic constants ──────────────────────────────────────────────────────

_ROS_VESC_LEFT     = "/vesc/left/state"
_ROS_VESC_RIGHT    = "/vesc/right/state"
_ROS_VESC_COMBINED = "/vesc/combined"


def _extract_motor_payload(data: dict) -> dict:
    """
    Extract the required fields from a vesc_telemetry_node per-motor JSON dict.

    Upstream JSON keys (from vesc_telemetry_node._state_to_dict):
      rpm, current_a, voltage_v, temp_fet_c, duty_cycle, fault_code, ...

    Returns a dashboard-friendly payload with a stable key set.
    """
    return {
        "rpm":           int(data["rpm"]),
        "current_a":     float(data["current_a"]),
        "voltage_v":     float(data["voltage_v"]),
        "temperature_c": float(data["temp_fet_c"]),
        "duty_cycle":    float(data["duty_cycle"]),
        "fault_code":    int(data["fault_code"]),
        "ts":            float(data.get("stamp", time.time())),
    }


def _extract_combined_payload(data: dict) -> dict:
    """
    Extract fields from the /vesc/combined JSON dict.

    Upstream keys: voltage_v, total_current_a, left_rpm, right_rpm,
                   left_alive, right_alive, stamp
    """
    return {
        "voltage_v":       float(data["voltage_v"]),
        "total_current_a": float(data["total_current_a"]),
        "left_rpm":        int(data["left_rpm"]),
        "right_rpm":       int(data["right_rpm"]),
        "left_alive":      bool(data["left_alive"]),
        "right_alive":     bool(data["right_alive"]),
        "ts":              float(data.get("stamp", time.time())),
    }


class VescMqttRelayNode(Node):
    """
    Subscribes to VESC ROS2 topics and relays telemetry to MQTT.

    Rate limiting: each per-motor topic maintains a last-publish timestamp;
    messages arriving faster than motor_rate_hz are silently dropped.
    The /vesc/combined topic is also rate-limited at the same rate.
    """

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("mqtt_host",         "localhost")
        self.declare_parameter("mqtt_port",         1883)
        self.declare_parameter("mqtt_keepalive",    60)
        self.declare_parameter("reconnect_delay_s", 5.0)
        self.declare_parameter("motor_rate_hz",     5.0)

        self._mqtt_host        = self.get_parameter("mqtt_host").value
        self._mqtt_port        = self.get_parameter("mqtt_port").value
        self._mqtt_keepalive   = self.get_parameter("mqtt_keepalive").value
        reconnect_delay        = self.get_parameter("reconnect_delay_s").value
        rate_hz                = max(0.1, float(self.get_parameter("motor_rate_hz").value))
        self._min_interval     = 1.0 / rate_hz

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

        # ── Rate-limit state (last publish epoch per MQTT topic) ──────────────
        self._last_pub: dict[str, float] = {
            _MQTT_VESC_LEFT:     0.0,
            _MQTT_VESC_RIGHT:    0.0,
            _MQTT_VESC_COMBINED: 0.0,
        }

        # ── Stats ─────────────────────────────────────────────────────────────
        self._rx_count:    dict[str, int] = {
            _MQTT_VESC_LEFT: 0, _MQTT_VESC_RIGHT: 0, _MQTT_VESC_COMBINED: 0
        }
        self._pub_count:   dict[str, int] = {
            _MQTT_VESC_LEFT: 0, _MQTT_VESC_RIGHT: 0, _MQTT_VESC_COMBINED: 0
        }
        self._drop_count:  dict[str, int] = {
            _MQTT_VESC_LEFT: 0, _MQTT_VESC_RIGHT: 0, _MQTT_VESC_COMBINED: 0
        }
        self._err_count:   dict[str, int] = {
            _MQTT_VESC_LEFT: 0, _MQTT_VESC_RIGHT: 0, _MQTT_VESC_COMBINED: 0
        }
        self._mqtt_connected = False

        # ── ROS2 subscriptions ────────────────────────────────────────────────
        self.create_subscription(
            String, _ROS_VESC_LEFT,
            lambda msg: self._on_vesc(msg, _MQTT_VESC_LEFT, _extract_motor_payload),
            10,
        )
        self.create_subscription(
            String, _ROS_VESC_RIGHT,
            lambda msg: self._on_vesc(msg, _MQTT_VESC_RIGHT, _extract_motor_payload),
            10,
        )
        self.create_subscription(
            String, _ROS_VESC_COMBINED,
            lambda msg: self._on_vesc(msg, _MQTT_VESC_COMBINED, _extract_combined_payload),
            10,
        )

        # ── MQTT client ───────────────────────────────────────────────────────
        self._mqtt_client: "mqtt.Client | None" = None
        if _MQTT_OK:
            self._mqtt_client = mqtt.Client(
                client_id="saltybot-vesc-mqtt-relay", clean_session=True
            )
            self._mqtt_client.on_connect    = self._on_mqtt_connect
            self._mqtt_client.on_disconnect = self._on_mqtt_disconnect
            self._mqtt_client.reconnect_delay_set(
                min_delay=int(reconnect_delay),
                max_delay=int(reconnect_delay) * 4,
            )
            self._mqtt_connect()

        self.get_logger().info(
            "VESC→MQTT relay started — broker=%s:%d  rate=%.1f Hz",
            self._mqtt_host, self._mqtt_port, rate_hz,
        )

    # ── MQTT connection ───────────────────────────────────────────────────────

    def _mqtt_connect(self) -> None:
        try:
            self._mqtt_client.connect_async(
                self._mqtt_host, self._mqtt_port,
                keepalive=self._mqtt_keepalive,
            )
            self._mqtt_client.loop_start()
        except Exception as exc:
            self.get_logger().warning("MQTT initial connect error: %s", str(exc))

    def _on_mqtt_connect(self, client, userdata, flags, rc) -> None:
        if rc == 0:
            self._mqtt_connected = True
            self.get_logger().info(
                "MQTT connected to %s:%d", self._mqtt_host, self._mqtt_port
            )
        else:
            self.get_logger().warning("MQTT connect failed rc=%d", rc)

    def _on_mqtt_disconnect(self, client, userdata, rc) -> None:
        self._mqtt_connected = False
        if rc != 0:
            self.get_logger().warning(
                "MQTT disconnected (rc=%d) — paho will reconnect", rc
            )

    # ── ROS2 subscriber callback ──────────────────────────────────────────────

    def _on_vesc(self, ros_msg: String, mqtt_topic: str, extractor) -> None:
        """
        Handle an incoming VESC ROS2 message.

        1. Parse JSON from the String payload.
        2. Check rate limit — drop if too soon.
        3. Extract dashboard fields.
        4. Publish to MQTT.
        """
        self._rx_count[mqtt_topic] += 1

        # Rate limit
        now = time.monotonic()
        if now - self._last_pub[mqtt_topic] < self._min_interval:
            self._drop_count[mqtt_topic] += 1
            return

        # Parse
        try:
            data = json.loads(ros_msg.data)
        except (json.JSONDecodeError, UnicodeDecodeError) as exc:
            self._err_count[mqtt_topic] += 1
            self.get_logger().debug("JSON error on %s: %s", mqtt_topic, exc)
            return

        # Extract
        try:
            payload = extractor(data)
        except (KeyError, TypeError, ValueError) as exc:
            self._err_count[mqtt_topic] += 1
            self.get_logger().debug("Payload error on %s: %s — %s", mqtt_topic, exc, data)
            return

        # Publish
        if self._mqtt_client is not None and self._mqtt_connected:
            try:
                self._mqtt_client.publish(
                    mqtt_topic,
                    json.dumps(payload, separators=(",", ":")),
                    qos=0,
                    retain=False,
                )
                self._last_pub[mqtt_topic] = now
                self._pub_count[mqtt_topic] += 1
            except Exception as exc:
                self._err_count[mqtt_topic] += 1
                self.get_logger().debug("MQTT publish error on %s: %s", mqtt_topic, exc)
        else:
            # Still update last_pub to avoid log spam when broker is down
            self._last_pub[mqtt_topic] = now
            self.get_logger().debug("MQTT not connected — dropped %s", mqtt_topic)

    # ── Cleanup ───────────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        if self._mqtt_client is not None:
            self._mqtt_client.loop_stop()
            self._mqtt_client.disconnect()
        super().destroy_node()


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

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


if __name__ == "__main__":
    main()
