"""jetson_stats.py — Jetson Orin hardware stats collection (Issue #631).

No ROS2 dependencies — pure Python, fully unit-testable.

Design
──────
JetsonStatsReader.read() returns a JetsonStats snapshot.  The reader
tries three backends in priority order:

  1. jtop  (jetson_stats pip package) — richest data, low overhead
  2. tegrastats subprocess           — always present on JetPack, parseable
  3. MockReader (tests / dev)        — deterministic values, no hardware

AlertChecker.check(stats, thresholds) returns a list of Alert objects
describing any threshold violations.  All thresholds are per-dataclass
config so they're easy to unit-test.

Tegrastats sample line (Jetson Orin NX, JetPack 6):
    RAM 1234/7772MB (lfb 2x2MB) SWAP 0/3886MB (cached 0MB)
    CPU [12%@1190,5%@729,...] EMC_FREQ 4%@2133 GPC_FREQ 4%@918
    AO@35C CPU@39.6C Tboard@32C Tdiode@35.3C GPU@37.6C tj@41.3C
    VDD_GPU_SOC 1866mW/1866mW VDD_CPU_CV 700mW/682mW VDD_IN 4002mW/3996mW
"""

from __future__ import annotations

import re
import shutil
import subprocess
import time
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Tuple

# psutil is available everywhere; used for disk stats fallback
try:
    import psutil as _psutil
    HAS_PSUTIL = True
except ImportError:
    HAS_PSUTIL = False

# jtop (jetson_stats) is Jetson-only
try:
    from jtop import jtop as _jtop
    HAS_JTOP = True
except ImportError:
    HAS_JTOP = False


# ── Data model ────────────────────────────────────────────────────────────────

@dataclass
class CpuCore:
    """Stats for one CPU core."""
    index:     int
    usage_pct: float   # 0–100
    freq_mhz:  int     = 0


@dataclass
class JetsonStats:
    """Snapshot of all monitored Jetson Orin metrics."""

    timestamp: float = field(default_factory=time.time)

    # CPU
    cpu_cores:   List[CpuCore] = field(default_factory=list)
    cpu_avg_pct: float = 0.0   # mean across online cores

    # GPU
    gpu_pct:      float = 0.0
    gpu_freq_mhz: int   = 0

    # RAM
    ram_used_mb:  int   = 0
    ram_total_mb: int   = 0
    ram_pct:      float = 0.0

    # SWAP
    swap_used_mb:  int = 0
    swap_total_mb: int = 0

    # Temperatures (name → °C)
    temperatures: Dict[str, float] = field(default_factory=dict)
    temp_max_c:   float = 0.0    # hottest sensor

    # Fan
    fan_pct: float = 0.0   # 0–100 (0 if unavailable)

    # Power (mW)
    power_total_mw: int = 0
    power_cpu_mw:   int = 0
    power_gpu_mw:   int = 0

    # Disk (/, primary partition)
    disk_used_gb:  float = 0.0
    disk_total_gb: float = 0.0
    disk_pct:      float = 0.0

    # Source identification
    source: str = "unknown"

    # ── Derived helpers ────────────────────────────────────────────────────

    def to_dict(self) -> dict:
        """Return JSON-serialisable dict (cpu_cores as plain dicts)."""
        d = asdict(self)
        # dataclasses.asdict already recurses into nested dataclasses
        return d

    @property
    def cpu_count(self) -> int:
        return len(self.cpu_cores)

    @property
    def ram_free_mb(self) -> int:
        return max(0, self.ram_total_mb - self.ram_used_mb)

    @property
    def disk_free_gb(self) -> float:
        return max(0.0, self.disk_total_gb - self.disk_used_gb)


# ── Alert model ───────────────────────────────────────────────────────────────

@dataclass
class AlertThresholds:
    """Configurable alert thresholds for Jetson stats (Issue #631 spec)."""
    cpu_pct:   float = 85.0    # alert when CPU average exceeds this
    gpu_pct:   float = 85.0    # alert when GPU usage exceeds this
    temp_c:    float = 80.0    # alert when any sensor exceeds this
    disk_pct:  float = 90.0    # alert when disk usage exceeds this
    ram_pct:   float = 90.0    # alert when RAM usage exceeds this
    power_mw:  float = 30_000  # alert when total power exceeds 30 W


@dataclass
class Alert:
    """A single threshold-violation alert."""
    component: str            # e.g. "cpu", "gpu", "temp_CPU", "disk"
    message:   str
    level:     int            # 1=WARN, 2=ERROR (matches DiagnosticStatus levels)
    value:     float = 0.0    # observed value
    threshold: float = 0.0    # configured threshold

    WARN  = 1
    ERROR = 2

    def is_error(self) -> bool:
        return self.level == self.ERROR

    def is_warn(self) -> bool:
        return self.level == self.WARN


class AlertChecker:
    """Check a JetsonStats snapshot against AlertThresholds.

    Returns a list of Alert objects (empty = all OK).
    Levels:
        WARN  → value is between threshold and threshold+10% (or single step)
        ERROR → value is above threshold+10% (or temp > threshold)
    """

    def check(
        self,
        stats:      JetsonStats,
        thresholds: AlertThresholds,
    ) -> List[Alert]:
        alerts: List[Alert] = []

        # CPU
        if stats.cpu_avg_pct >= thresholds.cpu_pct:
            level = (Alert.ERROR if stats.cpu_avg_pct >= thresholds.cpu_pct + 10
                     else Alert.WARN)
            alerts.append(Alert(
                component = "cpu",
                message   = (f"CPU avg {stats.cpu_avg_pct:.1f}% "
                             f">= {thresholds.cpu_pct:.0f}%"),
                level     = level,
                value     = stats.cpu_avg_pct,
                threshold = thresholds.cpu_pct,
            ))

        # GPU
        if stats.gpu_pct >= thresholds.gpu_pct:
            level = (Alert.ERROR if stats.gpu_pct >= thresholds.gpu_pct + 10
                     else Alert.WARN)
            alerts.append(Alert(
                component = "gpu",
                message   = (f"GPU {stats.gpu_pct:.1f}% "
                             f">= {thresholds.gpu_pct:.0f}%"),
                level     = level,
                value     = stats.gpu_pct,
                threshold = thresholds.gpu_pct,
            ))

        # Temperatures — check every sensor
        for sensor, temp_c in stats.temperatures.items():
            if temp_c >= thresholds.temp_c:
                level = (Alert.ERROR if temp_c >= thresholds.temp_c + 10
                         else Alert.WARN)
                alerts.append(Alert(
                    component = f"temp_{sensor}",
                    message   = (f"{sensor} {temp_c:.1f}°C "
                                 f">= {thresholds.temp_c:.0f}°C"),
                    level     = level,
                    value     = temp_c,
                    threshold = thresholds.temp_c,
                ))

        # Disk
        if stats.disk_pct >= thresholds.disk_pct:
            level = (Alert.ERROR if stats.disk_pct >= thresholds.disk_pct + 5
                     else Alert.WARN)
            alerts.append(Alert(
                component = "disk",
                message   = (f"Disk {stats.disk_pct:.1f}% "
                             f">= {thresholds.disk_pct:.0f}%"),
                level     = level,
                value     = stats.disk_pct,
                threshold = thresholds.disk_pct,
            ))

        # RAM
        if stats.ram_pct >= thresholds.ram_pct:
            level = (Alert.ERROR if stats.ram_pct >= thresholds.ram_pct + 5
                     else Alert.WARN)
            alerts.append(Alert(
                component = "ram",
                message   = (f"RAM {stats.ram_pct:.1f}% "
                             f">= {thresholds.ram_pct:.0f}%"),
                level     = level,
                value     = stats.ram_pct,
                threshold = thresholds.ram_pct,
            ))

        # Power
        if stats.power_total_mw >= thresholds.power_mw:
            alerts.append(Alert(
                component = "power",
                message   = (f"Power {stats.power_total_mw/1000:.1f}W "
                             f">= {thresholds.power_mw/1000:.0f}W"),
                level     = Alert.WARN,
                value     = stats.power_total_mw,
                threshold = thresholds.power_mw,
            ))

        return alerts


# ── Tegrastats parser ─────────────────────────────────────────────────────────

class TegrastatsParser:
    """Parse a single tegrastats output line into a JetsonStats snapshot."""

    # RAM 1234/7772MB (lfb …)
    _RE_RAM   = re.compile(r"RAM\s+(\d+)/(\d+)MB")
    # SWAP 0/3886MB
    _RE_SWAP  = re.compile(r"SWAP\s+(\d+)/(\d+)MB")
    # CPU [12%@1190,5%@729,…]
    _RE_CPU   = re.compile(r"CPU\s+\[([^\]]+)\]")
    # GPC_FREQ 4%@918  (GPU usage + freq)
    _RE_GPU   = re.compile(r"GPC_FREQ\s+(\d+)%@(\d+)")
    # SensorName@Temp  e.g. CPU@39.6C  GPU@37.6C  tj@41.3C
    _RE_TEMP  = re.compile(r"(\w+)@(\d+\.?\d*)C")
    # VDD_IN 4002mW/4002mW  (total board power: instant/average)
    _RE_VDD   = re.compile(r"VDD_IN\s+(\d+)mW/(\d+)mW")
    # VDD_CPU_CV 700mW
    _RE_VCPU  = re.compile(r"VDD_CPU_CV\s+(\d+)mW/(\d+)mW")
    # VDD_GPU_SOC 1866mW
    _RE_VGPU  = re.compile(r"VDD_GPU_SOC\s+(\d+)mW/(\d+)mW")

    def parse(self, line: str) -> JetsonStats:
        """Parse one tegrastats text line.  Returns JetsonStats."""
        stats = JetsonStats(source="tegrastats")

        # RAM
        m = self._RE_RAM.search(line)
        if m:
            stats.ram_used_mb  = int(m.group(1))
            stats.ram_total_mb = int(m.group(2))
            if stats.ram_total_mb > 0:
                stats.ram_pct = stats.ram_used_mb / stats.ram_total_mb * 100.0

        # SWAP
        m = self._RE_SWAP.search(line)
        if m:
            stats.swap_used_mb  = int(m.group(1))
            stats.swap_total_mb = int(m.group(2))

        # CPU cores
        m = self._RE_CPU.search(line)
        if m:
            cores: List[CpuCore] = []
            for idx, token in enumerate(m.group(1).split(",")):
                parts = token.split("@")
                pct_str = parts[0].rstrip("%")
                freq    = int(parts[1]) if len(parts) > 1 else 0
                try:
                    cores.append(CpuCore(index=idx,
                                         usage_pct=float(pct_str),
                                         freq_mhz=freq))
                except ValueError:
                    pass
            stats.cpu_cores   = cores
            if cores:
                stats.cpu_avg_pct = sum(c.usage_pct for c in cores) / len(cores)

        # GPU
        m = self._RE_GPU.search(line)
        if m:
            stats.gpu_pct      = float(m.group(1))
            stats.gpu_freq_mhz = int(m.group(2))

        # Temperatures
        temps: Dict[str, float] = {}
        for m in self._RE_TEMP.finditer(line):
            name, val = m.group(1), float(m.group(2))
            # Skip tokens that are part of CPU/freq strings like "1190"
            if name.isdigit():
                continue
            temps[name] = val
        stats.temperatures = temps
        stats.temp_max_c   = max(temps.values(), default=0.0)

        # Power
        m = self._RE_VDD.search(line)
        if m:
            stats.power_total_mw = int(m.group(1))
        m = self._RE_VCPU.search(line)
        if m:
            stats.power_cpu_mw = int(m.group(1))
        m = self._RE_VGPU.search(line)
        if m:
            stats.power_gpu_mw = int(m.group(1))

        # Disk — tegrastats does not include disk; use psutil fallback
        stats.disk_used_gb, stats.disk_total_gb, stats.disk_pct = \
            _read_disk_psutil()

        stats.timestamp = time.time()
        return stats


# ── Disk helper ───────────────────────────────────────────────────────────────

def _read_disk_psutil(path: str = "/") -> Tuple[float, float, float]:
    """Return (used_gb, total_gb, pct) for the given path."""
    if HAS_PSUTIL:
        try:
            u = _psutil.disk_usage(path)
            used_gb  = u.used  / 1024**3
            total_gb = u.total / 1024**3
            pct      = u.percent
            return used_gb, total_gb, pct
        except OSError:
            pass
    # shutil fallback
    try:
        u = shutil.disk_usage(path)
        total_gb = u.total / 1024**3
        used_gb  = (u.total - u.free) / 1024**3
        pct      = used_gb / total_gb * 100.0 if total_gb else 0.0
        return used_gb, total_gb, pct
    except OSError:
        return 0.0, 0.0, 0.0


# ── Stats readers ─────────────────────────────────────────────────────────────

class TegrastatsReader:
    """Reads one snapshot from tegrastats subprocess."""

    _TEGRASTATS_BIN = "tegrastats"

    def __init__(self, interval_ms: int = 200) -> None:
        self._interval_ms = interval_ms
        self._parser      = TegrastatsParser()

    def is_available(self) -> bool:
        return shutil.which(self._TEGRASTATS_BIN) is not None

    def read(self) -> Optional[JetsonStats]:
        """Launch tegrastats, capture one line, parse it."""
        if not self.is_available():
            return None
        try:
            # --interval N: emit every N ms; capture 1 line then kill
            result = subprocess.run(
                [self._TEGRASTATS_BIN, f"--interval", str(self._interval_ms)],
                capture_output=True,
                text=True,
                timeout=3.0,
            )
            # tegrastats may not exit on its own; take first non-empty line
            for line in result.stdout.splitlines():
                line = line.strip()
                if line:
                    return self._parser.parse(line)
        except (subprocess.TimeoutExpired, subprocess.SubprocessError,
                FileNotFoundError, OSError):
            pass
        return None


class JtopReader:
    """Reads one snapshot via the jtop Python library (jetson_stats)."""

    def is_available(self) -> bool:
        return HAS_JTOP

    def read(self) -> Optional[JetsonStats]:
        if not HAS_JTOP:
            return None
        try:
            with _jtop() as jetson:
                stats = JetsonStats(source="jtop")

                # CPU
                if hasattr(jetson, "cpu"):
                    cores = []
                    cpu_data = jetson.cpu
                    # jetson_stats 4.x: cpu is a dict with keys 'CPU1', 'CPU2', ...
                    if isinstance(cpu_data, dict):
                        for i, (k, v) in enumerate(sorted(cpu_data.items())):
                            if isinstance(v, dict):
                                pct  = float(v.get("val", 0) or 0)
                                freq = int(v.get("frq", 0) or 0)
                            else:
                                pct, freq = 0.0, 0
                            cores.append(CpuCore(index=i, usage_pct=pct,
                                                  freq_mhz=freq))
                    stats.cpu_cores   = cores
                    stats.cpu_avg_pct = (sum(c.usage_pct for c in cores)
                                         / len(cores)) if cores else 0.0

                # GPU
                if hasattr(jetson, "gpu"):
                    gpu_d = jetson.gpu
                    if isinstance(gpu_d, dict):
                        stats.gpu_pct      = float(gpu_d.get("val", 0) or 0)
                        stats.gpu_freq_mhz = int(gpu_d.get("frq", 0) or 0)

                # RAM
                if hasattr(jetson, "memory"):
                    mem = jetson.memory
                    if isinstance(mem, dict):
                        ram = mem.get("RAM", {})
                        if isinstance(ram, dict):
                            stats.ram_used_mb  = int(ram.get("used", 0) or 0)
                            stats.ram_total_mb = int(ram.get("tot",  0) or 0)
                        swap = mem.get("SWAP", {})
                        if isinstance(swap, dict):
                            stats.swap_used_mb  = int(swap.get("used", 0) or 0)
                            stats.swap_total_mb = int(swap.get("tot",  0) or 0)
                    if stats.ram_total_mb > 0:
                        stats.ram_pct = (stats.ram_used_mb /
                                         stats.ram_total_mb * 100.0)

                # Temperatures
                if hasattr(jetson, "temperature"):
                    temps = {}
                    for k, v in jetson.temperature.items():
                        if isinstance(v, dict):
                            val = v.get("temp", None)
                        else:
                            val = v
                        if val is not None:
                            try:
                                temps[str(k)] = float(val)
                            except (TypeError, ValueError):
                                pass
                    stats.temperatures = temps
                    stats.temp_max_c   = max(temps.values(), default=0.0)

                # Fan
                if hasattr(jetson, "fan"):
                    fan_d = jetson.fan
                    if isinstance(fan_d, dict):
                        stats.fan_pct = float(fan_d.get("speed", 0) or 0)

                # Power
                if hasattr(jetson, "power"):
                    pwr = jetson.power
                    if isinstance(pwr, dict):
                        tot = pwr.get("tot", {})
                        if isinstance(tot, dict):
                            stats.power_total_mw = int(
                                (tot.get("power", 0) or 0) * 1000)

                # Disk
                stats.disk_used_gb, stats.disk_total_gb, stats.disk_pct = \
                    _read_disk_psutil()

                stats.timestamp = time.time()
                return stats
        except Exception:
            return None


class MockReader:
    """Returns deterministic stats for testing and development."""

    def __init__(self, **overrides) -> None:
        self._overrides = overrides

    def is_available(self) -> bool:
        return True

    def read(self) -> JetsonStats:
        cores = [CpuCore(index=i, usage_pct=20.0 + i * 2, freq_mhz=1200)
                 for i in range(4)]
        s = JetsonStats(
            cpu_cores    = cores,
            cpu_avg_pct  = sum(c.usage_pct for c in cores) / len(cores),
            gpu_pct      = 30.0,
            gpu_freq_mhz = 918,
            ram_used_mb  = 3000,
            ram_total_mb = 7772,
            ram_pct      = 3000 / 7772 * 100.0,
            swap_used_mb = 0,
            swap_total_mb= 3886,
            temperatures = {"CPU": 40.0, "GPU": 38.0, "AO": 35.0, "tj": 42.0},
            temp_max_c   = 42.0,
            fan_pct      = 30.0,
            power_total_mw = 8000,
            power_cpu_mw   = 2000,
            power_gpu_mw   = 1500,
            disk_used_gb   = 50.0,
            disk_total_gb  = 500.0,
            disk_pct       = 10.0,
            source         = "mock",
        )
        for k, v in self._overrides.items():
            if hasattr(s, k):
                setattr(s, k, v)
        return s


# ── Auto-selecting reader factory ─────────────────────────────────────────────

def make_reader(prefer: str = "auto"):
    """Return the best available reader.

    prefer: "jtop" | "tegrastats" | "mock" | "auto"
    """
    if prefer == "mock":
        return MockReader()
    if prefer == "jtop" or (prefer == "auto" and HAS_JTOP):
        r = JtopReader()
        if r.is_available():
            return r
    if prefer == "tegrastats" or prefer == "auto":
        r = TegrastatsReader()
        if r.is_available():
            return r
    # Fallback: mock (dev machine without Jetson hw)
    return MockReader()
