#!/usr/bin/env python3
"""sensor_health_node.py — ROS2 sensor topic health monitor (Issue #566).

Monitors all sensor topics for staleness. Each sensor is checked against
configurable WARN and ERROR timeouts. Results are published as a ROS2
DiagnosticArray and as an MQTT JSON summary.

Monitored topics (configurable via sensor_health_params.yaml):
  /camera/color/image_raw             — RealSense colour stream
  /camera/depth/image_rect_raw        — RealSense depth stream
  /camera/color/camera_info           — RealSense camera info
  /scan                               — LiDAR 2-D scan
  /imu/data                           — IMU (BNO055 via JLink)
  /saltybot/uwb/range                 — UWB ranging
  /saltybot/battery                   — Battery telemetry (JSON string)
  /saltybot/motor_daemon/status       — Motor daemon status

Published topics:
  /saltybot/diagnostics   (diagnostic_msgs/DiagnosticArray)  — full per-sensor status
  /saltybot/sensor_health (std_msgs/String)                   — JSON summary

MQTT:
  Topic: saltybot/health
  Payload: JSON {timestamp, overall, sensors:{name: {status, age_s, hz}}}

Parameters (config/sensor_health_params.yaml):
  check_hz             1.0     Health check publish rate
  mqtt_broker          localhost
  mqtt_port            1883
  mqtt_topic           saltybot/health
  mqtt_enabled         true
  sensors              list of {topic, name, warn_s, error_s, critical}
"""

from __future__ import annotations

import json
import time
import threading
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    DurabilityPolicy, HistoryPolicy, QoSProfile, ReliabilityPolicy
)

from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue
from sensor_msgs.msg import CameraInfo, Image, Imu, LaserScan
from std_msgs.msg import String


# ── Diagnostic level aliases ───────────────────────────────────────────────────

OK    = DiagnosticStatus.OK     # 0
WARN  = DiagnosticStatus.WARN   # 1
ERROR = DiagnosticStatus.ERROR  # 2

_LEVEL_NAMES = {OK: "OK", WARN: "WARN", ERROR: "ERROR"}

# ── Default sensor configuration ──────────────────────────────────────────────

DEFAULT_SENSORS: List[dict] = [
    {"topic": "/camera/color/image_raw",        "name": "camera_color",     "warn_s": 1.0, "error_s": 3.0, "critical": True},
    {"topic": "/camera/depth/image_rect_raw",   "name": "camera_depth",     "warn_s": 1.0, "error_s": 3.0, "critical": True},
    {"topic": "/camera/color/camera_info",      "name": "camera_info",      "warn_s": 2.0, "error_s": 5.0, "critical": False},
    {"topic": "/scan",                          "name": "lidar",            "warn_s": 1.0, "error_s": 3.0, "critical": True},
    {"topic": "/imu/data",                      "name": "imu",              "warn_s": 0.5, "error_s": 2.0, "critical": True},
    {"topic": "/saltybot/uwb/range",            "name": "uwb",              "warn_s": 2.0, "error_s": 5.0, "critical": False},
    {"topic": "/saltybot/battery",              "name": "battery",          "warn_s": 3.0, "error_s": 8.0, "critical": False},
    {"topic": "/saltybot/motor_daemon/status",  "name": "motor_daemon",     "warn_s": 2.0, "error_s": 5.0, "critical": True},
]


# ── SensorWatcher ─────────────────────────────────────────────────────────────

class SensorWatcher:
    """Tracks message receipt timestamps and rate for a single topic.

    Thread-safe: callback may fire from any executor thread.
    """

    def __init__(self, topic: str, name: str,
                 warn_s: float, error_s: float, critical: bool) -> None:
        self.topic    = topic
        self.name     = name
        self.warn_s   = warn_s
        self.error_s  = error_s
        self.critical = critical

        self._last_rx:  float = 0.0   # monotonic; 0 = never received
        self._count:    int   = 0
        self._rate_hz:  float = 0.0
        self._rate_ts:  float = time.monotonic()
        self._rate_cnt: int   = 0
        self._lock = threading.Lock()

    # ── Callback (called from subscription) ────────────────────────────────

    def on_message(self, _msg) -> None:
        """Record receipt of any message on this topic."""
        now = time.monotonic()
        with self._lock:
            self._last_rx = now
            self._count  += 1
            self._rate_cnt += 1
            # Update rate estimate every ~2 s
            elapsed = now - self._rate_ts
            if elapsed >= 2.0:
                self._rate_hz  = self._rate_cnt / elapsed
                self._rate_cnt = 0
                self._rate_ts  = now

    # ── Status query ───────────────────────────────────────────────────────

    def age_s(self, now: Optional[float] = None) -> float:
        """Seconds since last message (∞ if never received)."""
        if now is None:
            now = time.monotonic()
        with self._lock:
            if self._last_rx == 0.0:
                return float("inf")
            return now - self._last_rx

    def rate_hz(self) -> float:
        with self._lock:
            return self._rate_hz

    def msg_count(self) -> int:
        with self._lock:
            return self._count

    def level(self, now: Optional[float] = None) -> int:
        age = self.age_s(now)
        if age >= self.error_s:
            return ERROR
        if age >= self.warn_s:
            return WARN
        return OK

    def diagnostic_status(self, now: Optional[float] = None) -> DiagnosticStatus:
        if now is None:
            now = time.monotonic()
        age   = self.age_s(now)
        lvl   = self.level(now)
        hz    = self.rate_hz()
        cnt   = self.msg_count()

        if age == float("inf"):
            msg = f"ERROR: no data received on {self.topic}"
        elif lvl == ERROR:
            msg = f"ERROR: stale {age:.1f}s (threshold {self.error_s:.1f}s)"
        elif lvl == WARN:
            msg = f"WARN: stale {age:.1f}s (threshold {self.warn_s:.1f}s)"
        else:
            msg = f"OK ({hz:.1f} Hz)"

        age_str = "inf" if age == float("inf") else f"{age:.2f}"
        status = DiagnosticStatus()
        status.level   = lvl
        status.name    = self.name
        status.message = msg
        status.hardware_id = self.topic
        status.values  = [
            KeyValue(key="topic",    value=self.topic),
            KeyValue(key="age_s",    value=age_str),
            KeyValue(key="rate_hz",  value=f"{hz:.2f}"),
            KeyValue(key="count",    value=str(cnt)),
            KeyValue(key="warn_s",   value=str(self.warn_s)),
            KeyValue(key="error_s",  value=str(self.error_s)),
            KeyValue(key="critical", value=str(self.critical)),
        ]
        return status

    def summary_dict(self, now: Optional[float] = None) -> dict:
        if now is None:
            now = time.monotonic()
        age = self.age_s(now)
        return {
            "status":    _LEVEL_NAMES[self.level(now)],
            "age_s":     round(age, 2) if age != float("inf") else None,
            "rate_hz":   round(self.rate_hz(), 2),
            "count":     self.msg_count(),
            "critical":  self.critical,
        }


# ── SensorHealthNode ───────────────────────────────────────────────────────────

class SensorHealthNode(Node):
    """Monitor all sensor topics; publish DiagnosticArray + MQTT JSON."""

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

        # ── Parameters ─────────────────────────────────────────────────────
        self.declare_parameter("check_hz",    1.0)
        self.declare_parameter("mqtt_broker", "localhost")
        self.declare_parameter("mqtt_port",   1883)
        self.declare_parameter("mqtt_topic",  "saltybot/health")
        self.declare_parameter("mqtt_enabled", True)

        check_hz         = self.get_parameter("check_hz").value
        self._mqtt_broker = self.get_parameter("mqtt_broker").value
        self._mqtt_port   = int(self.get_parameter("mqtt_port").value)
        self._mqtt_topic  = self.get_parameter("mqtt_topic").value
        mqtt_enabled      = self.get_parameter("mqtt_enabled").value

        # ── Build sensor watchers ───────────────────────────────────────────
        self._watchers: Dict[str, SensorWatcher] = {}
        for cfg in DEFAULT_SENSORS:
            w = SensorWatcher(
                topic    = cfg["topic"],
                name     = cfg["name"],
                warn_s   = float(cfg.get("warn_s",   1.0)),
                error_s  = float(cfg.get("error_s",  3.0)),
                critical = bool(cfg.get("critical",  False)),
            )
            self._watchers[cfg["name"]] = w

        # ── Subscriptions — one per sensor type ────────────────────────────
        # Best-effort QoS for sensor data (sensors may use BEST_EFFORT publishers)
        be_qos = QoSProfile(
            history=HistoryPolicy.KEEP_LAST, depth=1,
            reliability=ReliabilityPolicy.BEST_EFFORT,
            durability=DurabilityPolicy.VOLATILE,
        )
        rel_qos = QoSProfile(
            history=HistoryPolicy.KEEP_LAST, depth=5,
            reliability=ReliabilityPolicy.RELIABLE,
        )

        self._subscribe(Image,      "/camera/color/image_raw",       "camera_color",  be_qos)
        self._subscribe(Image,      "/camera/depth/image_rect_raw",  "camera_depth",  be_qos)
        self._subscribe(CameraInfo, "/camera/color/camera_info",     "camera_info",   be_qos)
        self._subscribe(LaserScan,  "/scan",                         "lidar",         be_qos)
        self._subscribe(Imu,        "/imu/data",                     "imu",           be_qos)
        self._subscribe(String,     "/saltybot/uwb/range",           "uwb",           rel_qos)
        self._subscribe(String,     "/saltybot/battery",             "battery",       rel_qos)
        self._subscribe(String,     "/saltybot/motor_daemon/status", "motor_daemon",  rel_qos)

        # ── Publishers ─────────────────────────────────────────────────────
        self._pub_diag   = self.create_publisher(
            DiagnosticArray, "/saltybot/diagnostics",   10)
        self._pub_health = self.create_publisher(
            String,          "/saltybot/sensor_health", 10)

        # ── MQTT ───────────────────────────────────────────────────────────
        self._mqtt_client = None
        if mqtt_enabled:
            self._connect_mqtt()

        # ── Timer ──────────────────────────────────────────────────────────
        self.create_timer(1.0 / max(0.1, check_hz), self._publish_diagnostics)

        sensor_list = ", ".join(self._watchers.keys())
        self.get_logger().info(
            f"[sensor_health] monitoring {len(self._watchers)} sensors: {sensor_list}"
        )

    # ── Subscription helper ────────────────────────────────────────────────

    def _subscribe(self, msg_type, topic: str, name: str, qos) -> None:
        if name not in self._watchers:
            return
        watcher = self._watchers[name]
        self.create_subscription(msg_type, topic, watcher.on_message, qos)

    # ── MQTT ───────────────────────────────────────────────────────────────

    def _connect_mqtt(self) -> None:
        try:
            import paho.mqtt.client as mqtt  # type: ignore
            self._mqtt_client = mqtt.Client(client_id="saltybot_sensor_health")
            self._mqtt_client.on_connect = self._on_mqtt_connect
            self._mqtt_client.connect_async(self._mqtt_broker, self._mqtt_port, 60)
            self._mqtt_client.loop_start()
        except ImportError:
            self.get_logger().warn(
                "[sensor_health] paho-mqtt not installed — MQTT publishing disabled"
            )
        except Exception as e:
            self.get_logger().warn(f"[sensor_health] MQTT connect failed: {e}")

    def _on_mqtt_connect(self, client, userdata, flags, rc) -> None:
        if rc == 0:
            self.get_logger().info(
                f"[sensor_health] MQTT connected to {self._mqtt_broker}:{self._mqtt_port}"
            )
        else:
            self.get_logger().warn(f"[sensor_health] MQTT connect rc={rc}")

    def _publish_mqtt(self, payload: str) -> None:
        if self._mqtt_client is None:
            return
        try:
            self._mqtt_client.publish(self._mqtt_topic, payload, qos=0, retain=True)
        except Exception as e:
            self.get_logger().warn(f"[sensor_health] MQTT publish error: {e}")

    # ── Diagnostic publisher ───────────────────────────────────────────────

    def _publish_diagnostics(self) -> None:
        now    = time.monotonic()
        wall   = time.time()

        # Build DiagnosticArray
        diag_array = DiagnosticArray()
        diag_array.header.stamp = self.get_clock().now().to_msg()

        sensor_summaries: dict = {}
        worst_level = OK

        for name, watcher in self._watchers.items():
            ds = watcher.diagnostic_status(now)
            diag_array.status.append(ds)

            if ds.level > worst_level:
                worst_level = ds.level

            sensor_summaries[name] = watcher.summary_dict(now)

        # Overall system status
        n_error   = sum(1 for w in self._watchers.values() if w.level(now) == ERROR)
        n_warn    = sum(1 for w in self._watchers.values() if w.level(now) == WARN)
        crit_err  = [n for n, w in self._watchers.items()
                     if w.critical and w.level(now) == ERROR]

        overall = "OK"
        if crit_err:
            overall = "ERROR"
        elif n_error > 0:
            overall = "ERROR"
        elif n_warn > 0:
            overall = "WARN"

        # Publish DiagnosticArray
        self._pub_diag.publish(diag_array)

        # JSON summary
        summary = {
            "timestamp":    wall,
            "overall":      overall,
            "n_ok":         len(self._watchers) - n_error - n_warn,
            "n_warn":       n_warn,
            "n_error":      n_error,
            "critical_err": crit_err,
            "sensors":      sensor_summaries,
        }
        payload = json.dumps(summary)
        self._pub_health.publish(String(data=payload))
        self._publish_mqtt(payload)

        # Log on transitions or errors
        if worst_level >= ERROR:
            self.get_logger().warn(
                f"[sensor_health] {overall}: {n_error} error(s), {n_warn} warn(s)"
                + (f" — critical: {crit_err}" if crit_err else "")
            )

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


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


if __name__ == "__main__":
    main()
