"""sysmon_node.py — System resource monitor for Jetson Orin.
Issue #355

Reads CPU, GPU, RAM, disk usage, and thermal temperatures, then
publishes a JSON string to /saltybot/system_resources at a configurable
rate (default 1 Hz).

All reads use /proc and /sys where available; GPU load falls back to -1.0
if the sysfs path is absent (non-Jetson host).

JSON payload
────────────
{
  "ts":              1234567890.123,   // epoch seconds (float)
  "cpu_percent":    [45.2, 32.1],     // per-core %; index 0 = aggregate
  "cpu_avg_percent": 38.6,            // mean of per-core values
  "ram_total_mb":   16384.0,
  "ram_used_mb":    4096.0,
  "ram_percent":    25.0,
  "disk_total_gb":  64.0,
  "disk_used_gb":   12.5,
  "disk_percent":   19.5,
  "gpu_percent":    42.0,             // -1.0 if unavailable
  "thermal":        {"CPU-therm": 47.5, "GPU-therm": 43.2}
}

Publications
────────────
  /saltybot/system_resources   std_msgs/String   JSON payload

Parameters
──────────
  publish_rate     (float, 1.0)          publish frequency (Hz)
  disk_path        (str,   "/")          path for statvfs disk usage
  gpu_sysfs_path   (str,   "/sys/devices/gpu.0/load")
  thermal_glob     (str,   "/sys/devices/virtual/thermal/thermal_zone*/temp")
  thermal_type_glob (str,  "/sys/devices/virtual/thermal/thermal_zone*/type")
  output_topic     (str,   "/saltybot/system_resources")
"""

from __future__ import annotations

import glob as _glob_mod
import json
import os
import threading
import time
from typing import Dict, List, Optional, Tuple

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


# ── Pure helpers ────────────────────────────────────────────────────────────────

def parse_proc_stat(content: str) -> List[List[int]]:
    """Parse /proc/stat and return per-cpu jiffies lists.

    Returns a list where index 0 = "cpu" (aggregate) and index 1+ = cpu0, cpu1, …
    Each entry is [user, nice, system, idle, iowait, irq, softirq, steal].
    """
    result: List[List[int]] = []
    for line in content.splitlines():
        if not line.startswith("cpu"):
            break
        parts = line.split()
        # parts[0] is "cpu" or "cpu0", parts[1:] are jiffie fields
        fields = [int(x) for x in parts[1:9]]  # take up to 8 fields
        while len(fields) < 8:
            fields.append(0)
        result.append(fields)
    return result


def cpu_percent_from_stats(
    prev: List[List[int]],
    curr: List[List[int]],
) -> List[float]:
    """Compute per-cpu busy percentage from two /proc/stat snapshots.

    Index 0 = aggregate (from "cpu" line), index 1+ = per-core.
    Returns empty list if inputs are incompatible.
    """
    if not prev or not curr or len(prev) != len(curr):
        return []

    result: List[float] = []
    for p, c in zip(prev, curr):
        idle_p = p[3] + p[4]   # idle + iowait
        idle_c = c[3] + c[4]
        total_p = sum(p)
        total_c = sum(c)
        d_total = total_c - total_p
        d_idle  = idle_c  - idle_p
        if d_total <= 0:
            result.append(0.0)
        else:
            pct = 100.0 * (d_total - d_idle) / d_total
            result.append(round(max(0.0, min(100.0, pct)), 2))
    return result


def parse_meminfo(content: str) -> Dict[str, int]:
    """Parse /proc/meminfo, return mapping of key → value in kB."""
    info: Dict[str, int] = {}
    for line in content.splitlines():
        if ":" not in line:
            continue
        key, _, rest = line.partition(":")
        value_str = rest.strip().split()[0] if rest.strip() else "0"
        try:
            info[key.strip()] = int(value_str)
        except ValueError:
            pass
    return info


def compute_ram_stats(meminfo: Dict[str, int]) -> Tuple[float, float, float]:
    """Return (total_mb, used_mb, percent) from parsed /proc/meminfo.

    Uses MemAvailable when present; falls back to MemFree.
    """
    total_kb    = meminfo.get("MemTotal", 0)
    avail_kb    = meminfo.get("MemAvailable", meminfo.get("MemFree", 0))
    used_kb     = max(0, total_kb - avail_kb)
    total_mb    = round(total_kb / 1024.0, 2)
    used_mb     = round(used_kb  / 1024.0, 2)
    if total_kb > 0:
        percent = round(100.0 * used_kb / total_kb, 2)
    else:
        percent = 0.0
    return total_mb, used_mb, percent


def read_disk_usage(path: str) -> Tuple[float, float, float]:
    """Return (total_gb, used_gb, percent) for the filesystem at *path*.

    Returns (-1.0, -1.0, -1.0) on error.
    """
    try:
        st = os.statvfs(path)
        total = st.f_frsize * st.f_blocks
        free  = st.f_frsize * st.f_bavail
        used  = total - free
        total_gb = round(total / (1024 ** 3), 3)
        used_gb  = round(used  / (1024 ** 3), 3)
        pct      = round(100.0 * used / total, 2) if total > 0 else 0.0
        return total_gb, used_gb, pct
    except Exception:
        return -1.0, -1.0, -1.0


def read_gpu_load(path: str) -> float:
    """Read Jetson GPU load from sysfs (0–100) or -1.0 if unavailable."""
    try:
        with open(path, "r") as fh:
            raw = fh.read().strip()
        # Some Jetson kernels report 0–1000 (per-mille), others 0–100
        val = int(raw)
        if val > 100:
            val = round(val / 10.0, 1)
        return float(max(0.0, min(100.0, val)))
    except Exception:
        return -1.0


def read_thermal_zones(
    temp_glob: str,
    type_glob: str,
) -> Dict[str, float]:
    """Glob thermal zone sysfs paths and return {zone_type: temp_C}.

    Falls back to numeric zone names ("zone0", "zone1", …) when type
    files are absent.
    """
    temp_paths = sorted(_glob_mod.glob(temp_glob))
    result: Dict[str, float] = {}

    for tp in temp_paths:
        # Derive zone directory from temp path: …/thermal_zoneN/temp
        zone_dir = os.path.dirname(tp)
        zone_index = os.path.basename(zone_dir).replace("thermal_zone", "zone")

        # Try to read the human-readable type
        type_path = os.path.join(zone_dir, "type")
        try:
            with open(type_path, "r") as fh:
                zone_name = fh.read().strip()
        except Exception:
            zone_name = zone_index

        try:
            with open(tp, "r") as fh:
                milli_c = int(fh.read().strip())
            result[zone_name] = round(milli_c / 1000.0, 1)
        except Exception:
            pass

    return result


# ── ROS2 node ───────────────────────────────────────────────────────────────────

class SysmonNode(Node):
    """Publishes Jetson system resource usage as a JSON String at a fixed rate."""

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

        self.declare_parameter("publish_rate",      1.0)
        self.declare_parameter("disk_path",         "/")
        self.declare_parameter("gpu_sysfs_path",    "/sys/devices/gpu.0/load")
        self.declare_parameter(
            "thermal_glob",
            "/sys/devices/virtual/thermal/thermal_zone*/temp"
        )
        self.declare_parameter(
            "thermal_type_glob",
            "/sys/devices/virtual/thermal/thermal_zone*/type"
        )
        self.declare_parameter("output_topic",      "/saltybot/system_resources")

        rate          = float(self.get_parameter("publish_rate").value)
        self._disk    = self.get_parameter("disk_path").value
        self._gpu_path = self.get_parameter("gpu_sysfs_path").value
        self._th_temp = self.get_parameter("thermal_glob").value
        self._th_type = self.get_parameter("thermal_type_glob").value
        topic         = self.get_parameter("output_topic").value

        # Injectable I/O functions for offline testing
        self._read_proc_stat  = self._default_read_proc_stat
        self._read_meminfo    = self._default_read_meminfo
        self._read_disk_usage = read_disk_usage
        self._read_gpu_load   = read_gpu_load
        self._read_thermal    = read_thermal_zones

        # CPU stats require two samples; prime with first read
        self._lock        = threading.Lock()
        self._prev_stat: Optional[List[List[int]]] = self._read_proc_stat()

        qos = QoSProfile(depth=10)
        self._pub   = self.create_publisher(String, topic, qos)
        self._timer = self.create_timer(1.0 / rate, self._tick)

        self.get_logger().info(
            f"SysmonNode ready — publishing {topic} at {rate:.1f} Hz"
        )

    # ── Default I/O readers ──────────────────────────────────────────────

    @staticmethod
    def _default_read_proc_stat() -> Optional[List[List[int]]]:
        try:
            with open("/proc/stat", "r") as fh:
                return parse_proc_stat(fh.read())
        except Exception:
            return None

    @staticmethod
    def _default_read_meminfo() -> str:
        try:
            with open("/proc/meminfo", "r") as fh:
                return fh.read()
        except Exception:
            return ""

    # ── Timer callback ───────────────────────────────────────────────────

    def _tick(self) -> None:
        ts = time.time()

        # CPU
        curr_stat = self._read_proc_stat()
        with self._lock:
            prev_stat = self._prev_stat
            self._prev_stat = curr_stat

        if prev_stat is not None and curr_stat is not None:
            cpu_pcts = cpu_percent_from_stats(prev_stat, curr_stat)
        else:
            cpu_pcts = []

        # Use per-core values (skip index 0 = aggregate) for avg
        per_core = cpu_pcts[1:] if len(cpu_pcts) > 1 else cpu_pcts
        cpu_avg  = round(sum(per_core) / len(per_core), 2) if per_core else 0.0

        # RAM
        meminfo = parse_meminfo(self._read_meminfo())
        ram_total, ram_used, ram_pct = compute_ram_stats(meminfo)

        # Disk
        disk_total, disk_used, disk_pct = self._read_disk_usage(self._disk)

        # GPU
        gpu_pct = self._read_gpu_load(self._gpu_path)

        # Thermal
        thermal = self._read_thermal(self._th_temp, self._th_type)

        payload = {
            "ts":              ts,
            "cpu_percent":    cpu_pcts,
            "cpu_avg_percent": cpu_avg,
            "ram_total_mb":   ram_total,
            "ram_used_mb":    ram_used,
            "ram_percent":    ram_pct,
            "disk_total_gb":  disk_total,
            "disk_used_gb":   disk_used,
            "disk_percent":   disk_pct,
            "gpu_percent":    gpu_pct,
            "thermal":        thermal,
        }

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

    # ── Public accessor ──────────────────────────────────────────────────

    @property
    def prev_stat(self) -> Optional[List[List[int]]]:
        with self._lock:
            return self._prev_stat


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