#!/usr/bin/env python3
"""Diagnostic self-test system for SaltyBot.

Performs startup hardware checks and continuous runtime monitoring.
Publishes /saltybot/diagnostics (DiagnosticArray), triggers TTS boot result,
face animation, and logs to /home/seb/saltybot-data/diagnostics/.
"""

import json
import time
import subprocess
from pathlib import Path
from datetime import datetime
from typing import Dict
import threading

import yaml
import psutil
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue
from std_msgs.msg import String


class DiagnosticsNode(Node):
    """ROS2 node for system diagnostics and self-test."""

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

        self.declare_parameter("enable_startup_check", True)
        self.declare_parameter("enable_runtime_monitoring", True)
        self.declare_parameter("monitoring_frequency", 1.0)
        self.declare_parameter("log_directory", "/home/seb/saltybot-data/diagnostics")
        self.declare_parameter("config_file", "diagnostic_checks.yaml")

        self.enable_startup_check = self.get_parameter("enable_startup_check").value
        self.enable_runtime_monitoring = self.get_parameter("enable_runtime_monitoring").value
        self.monitoring_freq = self.get_parameter("monitoring_frequency").value
        self.log_dir = Path(self.get_parameter("log_directory").value)
        config_file = self.get_parameter("config_file").value

        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.config = self._load_config(config_file)

        self.hardware_checks = {}
        self.runtime_metrics = {}
        self.startup_complete = False
        self.startup_time = time.time()

        self.pub_diagnostics = self.create_publisher(DiagnosticArray, "/saltybot/diagnostics", 1)
        self.pub_tts = self.create_publisher(String, "/saltybot/tts_say", 1)
        self.pub_face = self.create_publisher(String, "/saltybot/face/boot_animation", 1)

        if self.enable_startup_check:
            check_thread = threading.Thread(target=self._run_startup_checks)
            check_thread.daemon = True
            check_thread.start()

        if self.enable_runtime_monitoring:
            period = 1.0 / self.monitoring_freq
            self.timer = self.create_timer(period, self._runtime_check)

        self.get_logger().info(
            f"Diagnostics initialized. Startup: {self.enable_startup_check}, "
            f"Runtime: {self.enable_runtime_monitoring}@{self.monitoring_freq}Hz"
        )

    def _load_config(self, config_file: str) -> dict:
        """Load diagnostic configuration from YAML."""
        try:
            if not Path(config_file).exists():
                share_dir = Path(__file__).parent.parent / "config"
                config_file = str(share_dir / config_file)
            with open(config_file, "r") as f:
                return yaml.safe_load(f) or {}
        except Exception as e:
            self.get_logger().warn(f"Failed to load config: {e}")
            return {}

    def _run_startup_checks(self):
        """Run startup hardware checks."""
        try:
            self.get_logger().info("Starting hardware diagnostic checks...")
            
            # Run all checks
            self._check_rplidar()
            self._check_realsense()
            self._check_vesc()
            self._check_jabra()
            self._check_stm32()
            self._check_servos()
            self._check_wifi()
            self._check_gps()
            self._check_disk()
            self._check_ram()

            self._announce_boot_result()
            self._trigger_face_animation()
            self.startup_complete = True
            self._log_diagnostics("startup")
            
        except Exception as e:
            self.get_logger().error(f"Startup checks failed: {e}")

    def _check_rplidar(self):
        if Path("/dev/ttyUSB0").exists():
            self.hardware_checks["rplidar"] = ("OK", "RPLIDAR detected", {"port": "/dev/ttyUSB0"})
        else:
            self.hardware_checks["rplidar"] = ("ERROR", "RPLIDAR not found", {})

    def _check_realsense(self):
        try:
            result = subprocess.run(["lsusb"], capture_output=True, text=True, timeout=2)
            if "RealSense" in result.stdout or "Intel" in result.stdout:
                self.hardware_checks["realsense"] = ("OK", "RealSense detected", {})
            else:
                self.hardware_checks["realsense"] = ("WARN", "RealSense not in lsusb", {})
        except:
            self.hardware_checks["realsense"] = ("WARN", "lsusb check failed", {})

    def _check_vesc(self):
        vesc_found = any(Path(p).exists() for p in ["/dev/ttyUSB1", "/dev/ttyUSB2", "/dev/ttyACM0"])
        if vesc_found:
            self.hardware_checks["vesc"] = ("OK", "VESC detected", {})
        else:
            self.hardware_checks["vesc"] = ("ERROR", "VESC not found", {})

    def _check_jabra(self):
        try:
            result = subprocess.run(["arecord", "-l"], capture_output=True, text=True, timeout=2)
            if "Jabra" in result.stdout or "jabra" in result.stdout.lower():
                self.hardware_checks["jabra"] = ("OK", "Jabra microphone detected", {})
            else:
                self.hardware_checks["jabra"] = ("WARN", "Jabra not detected", {})
        except:
            self.hardware_checks["jabra"] = ("WARN", "Audio check failed", {})

    def _check_stm32(self):
<<<<<<< HEAD
        self.hardware_checks["stm32"] = ("OK", "ESP32 bridge online", {})
=======
        self.hardware_checks["stm32"] = ("OK", "ESP32-S3 bridge online", {})
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)

    def _check_servos(self):
        try:
            result = subprocess.run(["i2cdetect", "-y", "1"], capture_output=True, text=True, timeout=2)
            self.hardware_checks["servos"] = ("OK", "I2C servos available", {})
        except:
            self.hardware_checks["servos"] = ("WARN", "I2C check failed", {})

    def _check_wifi(self):
        try:
            result = subprocess.run(["iwconfig"], capture_output=True, text=True, timeout=2)
            if "ESSID" in result.stdout:
                self.hardware_checks["wifi"] = ("OK", "WiFi connected", {})
            else:
                self.hardware_checks["wifi"] = ("WARN", "WiFi not connected", {})
        except:
            self.hardware_checks["wifi"] = ("WARN", "WiFi check failed", {})

    def _check_gps(self):
        self.hardware_checks["gps"] = ("OK", "GPS module ready", {})

    def _check_disk(self):
        try:
            disk = psutil.disk_usage("/home/seb")
            percent = disk.percent
            if percent > 90:
                self.hardware_checks["disk"] = ("ERROR", f"Disk {percent}% full", {})
            elif percent > 80:
                self.hardware_checks["disk"] = ("WARN", f"Disk {percent}% used", {})
            else:
                self.hardware_checks["disk"] = ("OK", f"Disk OK {percent}% used", {})
        except Exception as e:
            self.hardware_checks["disk"] = ("WARN", f"Disk check failed: {e}", {})

    def _check_ram(self):
        try:
            mem = psutil.virtual_memory()
            percent = mem.percent
            if percent > 90:
                self.hardware_checks["ram"] = ("ERROR", f"RAM {percent}% used", {})
            elif percent > 80:
                self.hardware_checks["ram"] = ("WARN", f"RAM {percent}% used", {})
            else:
                self.hardware_checks["ram"] = ("OK", f"RAM OK {percent}% used", {})
        except Exception as e:
            self.hardware_checks["ram"] = ("WARN", f"RAM check failed: {e}", {})

    def _announce_boot_result(self):
        errors = [k for k, (s, _, _) in self.hardware_checks.items() if s == "ERROR"]
        if errors:
            msg = f"Boot complete with errors. {', '.join(errors)} offline."
        else:
            msg = "Boot complete. All systems online."
        self.pub_tts.publish(String(data=msg))

    def _trigger_face_animation(self):
        errors = [k for k, (s, _, _) in self.hardware_checks.items() if s == "ERROR"]
        animation = "boot_error" if errors else "boot_success"
        self.pub_face.publish(String(data=animation))

    def _runtime_check(self):
        if not self.startup_complete:
            return
        
        self._check_gpu_temp()
        self._check_network()
        self._publish_diagnostics()

    def _check_gpu_temp(self):
        try:
            result = subprocess.run(
                ["cat", "/sys/devices/virtual/thermal/thermal_zone0/temp"],
                capture_output=True, text=True, timeout=1
            )
            temp_c = int(result.stdout.strip()) / 1000.0
            status = "ERROR" if temp_c > 85 else "WARN" if temp_c > 80 else "OK"
            self.runtime_metrics["gpu_temp"] = (status, f"GPU {temp_c:.1f}C", {})
        except:
            pass

    def _check_network(self):
        try:
            result = subprocess.run(
                ["ping", "-c", "1", "-W", "1", "8.8.8.8"],
                capture_output=True, text=True, timeout=2
            )
            if result.returncode == 0:
                self.runtime_metrics["network"] = ("OK", "Network OK", {})
            else:
                self.runtime_metrics["network"] = ("WARN", "No connectivity", {})
        except:
            pass

    def _publish_diagnostics(self):
        array = DiagnosticArray()
        array.header.stamp = self.get_clock().now().to_msg()

        for name, (status, msg, details) in list(self.hardware_checks.items()) + list(self.runtime_metrics.items()):
            diag = DiagnosticStatus()
            diag.name = f"saltybot/{name}"
            level_map = {"OK": 0, "WARN": 1, "ERROR": 2}
            diag.level = level_map.get(status, 3)
            diag.message = msg
            
            for k, v in (details or {}).items():
                kv = KeyValue()
                kv.key = k
                kv.value = str(v)
                diag.values.append(kv)
            
            array.status.append(diag)

        self.pub_diagnostics.publish(array)

    def _log_diagnostics(self, check_type: str):
        try:
            log_data = {
                "timestamp": datetime.now().isoformat(),
                "check_type": check_type,
                "hardware_checks": {
                    name: {"status": s, "message": m, "details": d}
                    for name, (s, m, d) in self.hardware_checks.items()
                },
                "runtime_metrics": {
                    name: {"status": s, "message": m}
                    for name, (s, m, _) in self.runtime_metrics.items()
                },
            }
            filename = self.log_dir / f"diagnostics_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
            with open(filename, "w") as f:
                json.dump(log_data, f, indent=2)
            self.get_logger().info(f"Logged to {filename}")
        except Exception as e:
            self.get_logger().error(f"Failed to log: {e}")


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


if __name__ == "__main__":
    main()
