#!/usr/bin/env python3
"""Diagnostics aggregator — unified health dashboard node (Issue #658).

Subscribes to telemetry and diagnostics topics from all subsystems, aggregates
them into a single /saltybot/system_health JSON publish at 1 Hz.

Subscribed topics
-----------------
  /vesc/health                    (std_msgs/String)            — motor VESC health JSON
  /diagnostics                    (diagnostic_msgs/DiagnosticArray) — ROS diagnostics bus
  /saltybot/safety_zone/status    (std_msgs/String)            — safety / estop state
  /saltybot/pose/status           (std_msgs/String)            — IMU / pose estimate state
  /saltybot/uwb/status            (std_msgs/String)            — UWB positioning state

Published topics
----------------
  /saltybot/system_health  (std_msgs/String) — JSON health dashboard at 1 Hz

JSON schema for /saltybot/system_health
----------------------------------------
{
  "overall_status": "OK" | "WARN" | "ERROR" | "STALE",
  "uptime_s": <float>,
  "subsystems": {
    "motors":   { "status": ..., "message": ..., "age_s": ..., "previous_status": ... },
    "battery":  { ... },
    "imu":      { ... },
    "uwb":      { ... },
    "lidar":    { ... },
    "camera":   { ... },
    "can_bus":  { ... },
    "estop":    { ... }
  },
  "last_error": { "subsystem": ..., "message": ..., "timestamp": ... } | null,
  "recent_transitions": [ { "subsystem", "from_status", "to_status",
                             "message", "timestamp_iso" }, ... ],
  "timestamp": "<ISO-8601>"
}

Parameters
----------
  heartbeat_hz          float  1.0    Publish rate (Hz)
  stale_timeout_s       float  5.0    Seconds without update → STALE
  transition_log_size   int    50     Max recent transitions kept in memory
"""

import json
import time
from collections import deque
from datetime import datetime, timezone
from typing import Optional

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

from .subsystem import (
    SubsystemState,
    Transition,
    STATUS_OK,
    STATUS_WARN,
    STATUS_ERROR,
    STATUS_STALE,
    STATUS_UNKNOWN,
    worse,
    ros_level_to_status,
    SUBSYSTEM_NAMES as _SUBSYSTEM_NAMES,
    keyword_to_subsystem as _keyword_to_subsystem,
)


class DiagnosticsAggregatorNode(Node):
    """Unified health dashboard node."""

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

        # ----------------------------------------------------------------
        # Parameters
        # ----------------------------------------------------------------
        self.declare_parameter("heartbeat_hz",        1.0)
        self.declare_parameter("stale_timeout_s",     5.0)
        self.declare_parameter("transition_log_size", 50)

        hz              = float(self.get_parameter("heartbeat_hz").value)
        stale_timeout   = float(self.get_parameter("stale_timeout_s").value)
        log_size        = int(self.get_parameter("transition_log_size").value)

        # ----------------------------------------------------------------
        # Subsystem state table
        # ----------------------------------------------------------------
        self._subsystems: dict[str, SubsystemState] = {
            name: SubsystemState(name=name, stale_timeout_s=stale_timeout)
            for name in _SUBSYSTEM_NAMES
        }
        self._transitions: deque[Transition] = deque(maxlen=log_size)
        self._last_error: Optional[dict] = None
        self._start_mono = time.monotonic()

        # ----------------------------------------------------------------
        # Subscriptions
        # ----------------------------------------------------------------
        self.create_subscription(String,        "/vesc/health",
                                 self._on_vesc_health,       10)
        self.create_subscription(DiagnosticArray, "/diagnostics",
                                 self._on_diagnostics,       10)
        self.create_subscription(String,        "/saltybot/safety_zone/status",
                                 self._on_safety_zone,       10)
        self.create_subscription(String,        "/saltybot/pose/status",
                                 self._on_pose_status,       10)
        self.create_subscription(String,        "/saltybot/uwb/status",
                                 self._on_uwb_status,        10)

        # ----------------------------------------------------------------
        # Publisher
        # ----------------------------------------------------------------
        self._pub = self.create_publisher(String, "/saltybot/system_health", 1)

        # ----------------------------------------------------------------
        # 1 Hz heartbeat timer
        # ----------------------------------------------------------------
        self.create_timer(1.0 / max(0.1, hz), self._on_timer)

        self.get_logger().info(
            f"diagnostics_aggregator: {len(_SUBSYSTEM_NAMES)} subsystems, "
            f"heartbeat={hz} Hz, stale_timeout={stale_timeout} s"
        )

    # ----------------------------------------------------------------
    # Topic callbacks
    # ----------------------------------------------------------------

    def _on_vesc_health(self, msg: String) -> None:
        """Parse /vesc/health JSON → motors subsystem."""
        now = time.monotonic()
        try:
            d = json.loads(msg.data)
            fault = d.get("fault_code", 0)
            alive = d.get("alive", True)
            if not alive:
                status, message = STATUS_STALE, "VESC offline"
            elif fault != 0:
                status  = STATUS_ERROR
                message = f"VESC fault {d.get('fault_name', fault)}"
            else:
                status  = STATUS_OK
                message = (
                    f"RPM={d.get('rpm', '?')} "
                    f"I={d.get('current_a', '?')} A "
                    f"V={d.get('voltage_v', '?')} V"
                )
            self._update("motors", status, message, now)
        except Exception as exc:
            self.get_logger().debug(f"/vesc/health parse error: {exc}")

    def _on_diagnostics(self, msg: DiagnosticArray) -> None:
        """Fan /diagnostics entries out to per-subsystem states."""
        now = time.monotonic()
        # Accumulate worst status per subsystem across all entries in this array
        pending: dict[str, tuple[str, str]] = {}  # subsystem → (status, message)

        for ds in msg.status:
            subsystem = _keyword_to_subsystem(ds.name)
            if subsystem is None:
                continue
            status  = ros_level_to_status(ds.level)
            message = ds.message or ""
            if subsystem not in pending:
                pending[subsystem] = (status, message)
            else:
                prev_s, prev_m = pending[subsystem]
                new_worst = worse(status, prev_s)
                pending[subsystem] = (
                    new_worst,
                    message if new_worst == status else prev_m,
                )

        for subsystem, (status, message) in pending.items():
            self._update(subsystem, status, message, now)

    def _on_safety_zone(self, msg: String) -> None:
        """Parse /saltybot/safety_zone/status JSON → estop subsystem."""
        now = time.monotonic()
        try:
            d = json.loads(msg.data)
            triggered = d.get("estop_triggered", d.get("triggered", False))
            active    = d.get("safety_zone_active", d.get("active", True))
            if triggered:
                status, message = STATUS_ERROR, "E-stop triggered"
            elif not active:
                status, message = STATUS_WARN, "Safety zone inactive"
            else:
                status, message = STATUS_OK, "Safety zone active"
            self._update("estop", status, message, now)
        except Exception as exc:
            self.get_logger().debug(f"/saltybot/safety_zone/status parse error: {exc}")

    def _on_pose_status(self, msg: String) -> None:
        """Parse /saltybot/pose/status JSON → imu subsystem."""
        now = time.monotonic()
        try:
            d = json.loads(msg.data)
            status_str = d.get("status", "OK").upper()
            if status_str not in (STATUS_OK, STATUS_WARN, STATUS_ERROR):
                status_str = STATUS_WARN
            message = d.get("message", d.get("msg", ""))
            self._update("imu", status_str, message, now)
        except Exception as exc:
            self.get_logger().debug(f"/saltybot/pose/status parse error: {exc}")

    def _on_uwb_status(self, msg: String) -> None:
        """Parse /saltybot/uwb/status JSON → uwb subsystem."""
        now = time.monotonic()
        try:
            d = json.loads(msg.data)
            status_str = d.get("status", "OK").upper()
            if status_str not in (STATUS_OK, STATUS_WARN, STATUS_ERROR):
                status_str = STATUS_WARN
            message = d.get("message", d.get("msg", ""))
            self._update("uwb", status_str, message, now)
        except Exception as exc:
            self.get_logger().debug(f"/saltybot/uwb/status parse error: {exc}")

    # ----------------------------------------------------------------
    # Internal helpers
    # ----------------------------------------------------------------

    def _update(self, subsystem: str, status: str, message: str, now: float) -> None:
        """Update subsystem state and record any transition."""
        s = self._subsystems.get(subsystem)
        if s is None:
            return
        transition = s.update(status, message, now)
        if transition is not None:
            self._transitions.append(transition)
            self.get_logger().info(
                f"[{subsystem}] {transition.from_status} → {transition.to_status}: {message}"
            )
            if transition.to_status == STATUS_ERROR:
                self._last_error = {
                    "subsystem": subsystem,
                    "message":   message,
                    "timestamp": transition.timestamp_iso,
                }

    def _apply_stale_checks(self, now: float) -> None:
        for s in self._subsystems.values():
            transition = s.apply_stale_check(now)
            if transition is not None:
                self._transitions.append(transition)
                self.get_logger().warn(
                    f"[{transition.subsystem}] went STALE (no data)"
                )

    # ----------------------------------------------------------------
    # Heartbeat timer
    # ----------------------------------------------------------------

    def _on_timer(self) -> None:
        now = time.monotonic()
        self._apply_stale_checks(now)

        # Overall status = worst across all subsystems
        overall = STATUS_OK
        for s in self._subsystems.values():
            overall = worse(overall, s.status)
        # UNKNOWN does not degrade overall if at least one subsystem is known
        known = [s for s in self._subsystems.values() if s.status != STATUS_UNKNOWN]
        if not known:
            overall = STATUS_UNKNOWN

        uptime = now - self._start_mono

        payload = {
            "overall_status": overall,
            "uptime_s":       round(uptime, 1),
            "subsystems":     {
                name: s.to_dict(now)
                for name, s in self._subsystems.items()
            },
            "last_error": self._last_error,
            "recent_transitions": [
                {
                    "subsystem":   t.subsystem,
                    "from_status": t.from_status,
                    "to_status":   t.to_status,
                    "message":     t.message,
                    "timestamp":   t.timestamp_iso,
                }
                for t in list(self._transitions)[-10:]  # last 10 in the JSON
            ],
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

        self._pub.publish(String(data=json.dumps(payload)))


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


if __name__ == "__main__":
    main()
