#!/usr/bin/env python3
"""
Telemetry Observer - Real Physical State Reading

Reads actual hardware metrics for Se-aware QUBO routing:
- GPU: temp, power, utilization (nvidia-smi)
- CPU: temps (lm-sensors), load average
- Memory: used, available, swap
- IO: wait percentage, disk utilization

This is Lane B: "AI Lab as a physical plant" - real physics coupling.
"""

import subprocess
import re
import os
from typing import Dict, Optional, NamedTuple
from dataclasses import dataclass
from pathlib import Path


@dataclass
class GPUMetrics:
    """GPU telemetry from nvidia-smi."""
    temp_c: float           # Temperature in Celsius
    power_w: float          # Power draw in Watts
    power_limit_w: float    # Power limit in Watts
    utilization_pct: float  # GPU utilization %
    memory_used_mb: float   # Memory used in MB
    memory_total_mb: float  # Total memory in MB
    available: bool         # Whether GPU was readable


@dataclass
class CPUMetrics:
    """CPU telemetry."""
    load_1m: float          # 1-minute load average
    load_5m: float          # 5-minute load average
    load_15m: float         # 15-minute load average
    temp_c: Optional[float] # Temperature (if lm-sensors available)
    core_count: int         # Number of CPU cores


@dataclass
class MemoryMetrics:
    """Memory telemetry from /proc/meminfo."""
    total_mb: float
    available_mb: float
    used_mb: float
    swap_total_mb: float
    swap_used_mb: float
    pressure_pct: float     # Memory pressure (used/total)


@dataclass
class IOMetrics:
    """IO telemetry."""
    io_wait_pct: float      # IO wait from /proc/stat
    disk_util_pct: float    # Disk utilization (if iostat available)


@dataclass
class TelemetrySnapshot:
    """Complete telemetry snapshot at a point in time."""
    gpu: GPUMetrics
    cpu: CPUMetrics
    memory: MemoryMetrics
    io: IOMetrics
    timestamp: float        # Unix timestamp

    def to_dict(self) -> Dict:
        """Convert to dictionary for JSONL logging."""
        return {
            "timestamp": self.timestamp,
            "gpu": {
                "temp_c": self.gpu.temp_c,
                "power_w": self.gpu.power_w,
                "power_limit_w": self.gpu.power_limit_w,
                "utilization_pct": self.gpu.utilization_pct,
                "memory_used_mb": self.gpu.memory_used_mb,
                "memory_total_mb": self.gpu.memory_total_mb,
                "available": self.gpu.available,
            },
            "cpu": {
                "load_1m": self.cpu.load_1m,
                "load_5m": self.cpu.load_5m,
                "load_15m": self.cpu.load_15m,
                "temp_c": self.cpu.temp_c,
                "core_count": self.cpu.core_count,
            },
            "memory": {
                "total_mb": self.memory.total_mb,
                "available_mb": self.memory.available_mb,
                "used_mb": self.memory.used_mb,
                "swap_total_mb": self.memory.swap_total_mb,
                "swap_used_mb": self.memory.swap_used_mb,
                "pressure_pct": self.memory.pressure_pct,
            },
            "io": {
                "io_wait_pct": self.io.io_wait_pct,
                "disk_util_pct": self.io.disk_util_pct,
            },
        }


class TelemetryObserver:
    """
    Reads physical state from the system.

    This is the OBSERVE step in the control loop:
    OBSERVE -> COMPUTE Se -> DECIDE (QUBO) -> ACT
    """

    def __init__(self):
        self._has_nvidia_smi = self._check_nvidia_smi()
        self._has_lm_sensors = self._check_lm_sensors()
        self._cpu_count = os.cpu_count() or 1

    def _check_nvidia_smi(self) -> bool:
        """Check if nvidia-smi is available."""
        try:
            # Use -L to list GPUs (works on all driver versions)
            # Note: --version doesn't work on older drivers
            result = subprocess.run(
                ["nvidia-smi", "-L"],
                capture_output=True,
                timeout=5
            )
            return result.returncode == 0
        except (FileNotFoundError, subprocess.TimeoutExpired):
            return False

    def _check_lm_sensors(self) -> bool:
        """Check if lm-sensors is available."""
        try:
            result = subprocess.run(
                ["sensors", "--version"],
                capture_output=True,
                timeout=5
            )
            return result.returncode == 0
        except (FileNotFoundError, subprocess.TimeoutExpired):
            return False

    def read_gpu(self) -> GPUMetrics:
        """
        Read GPU metrics via nvidia-smi.

        Query format: temp, power, power_limit, utilization, memory_used, memory_total
        """
        if not self._has_nvidia_smi:
            return GPUMetrics(
                temp_c=0.0, power_w=0.0, power_limit_w=0.0,
                utilization_pct=0.0, memory_used_mb=0.0, memory_total_mb=0.0,
                available=False
            )

        try:
            result = subprocess.run(
                [
                    "nvidia-smi",
                    "--query-gpu=temperature.gpu,power.draw,power.limit,"
                    "utilization.gpu,memory.used,memory.total",
                    "--format=csv,noheader,nounits"
                ],
                capture_output=True,
                text=True,
                timeout=5
            )

            if result.returncode != 0:
                raise RuntimeError(f"nvidia-smi failed: {result.stderr}")

            # Parse CSV output: "45, 150.00, 300.00, 85, 8192, 24576"
            # Note: Some values may be "[N/A]" on certain GPUs/drivers
            values = result.stdout.strip().split(", ")
            if len(values) >= 6:
                def safe_float(v: str, default: float = 0.0) -> float:
                    try:
                        return float(v)
                    except ValueError:
                        return default  # Handle "[N/A]" or other non-numeric

                return GPUMetrics(
                    temp_c=safe_float(values[0]),
                    power_w=safe_float(values[1]),
                    power_limit_w=safe_float(values[2], 300.0),  # Default power limit
                    utilization_pct=safe_float(values[3]),
                    memory_used_mb=safe_float(values[4]),
                    memory_total_mb=safe_float(values[5]),
                    available=True
                )

        except (subprocess.TimeoutExpired, ValueError, IndexError) as e:
            pass

        return GPUMetrics(
            temp_c=0.0, power_w=0.0, power_limit_w=0.0,
            utilization_pct=0.0, memory_used_mb=0.0, memory_total_mb=0.0,
            available=False
        )

    def read_cpu(self) -> CPUMetrics:
        """
        Read CPU metrics: load average and optionally temperature.
        """
        # Load average from /proc/loadavg
        load_1m, load_5m, load_15m = 0.0, 0.0, 0.0
        try:
            with open("/proc/loadavg", "r") as f:
                parts = f.read().strip().split()
                load_1m = float(parts[0])
                load_5m = float(parts[1])
                load_15m = float(parts[2])
        except (IOError, IndexError, ValueError):
            pass

        # CPU temperature via lm-sensors (if available)
        temp_c = None
        if self._has_lm_sensors:
            try:
                result = subprocess.run(
                    ["sensors"],
                    capture_output=True,
                    text=True,
                    timeout=5
                )
                # Parse for patterns like "Core 0:        +45.0°C" or "Tctl:         +55.0°C"
                temps = re.findall(r'(?:Core \d+|Tctl|Package id \d+):\s*\+?([\d.]+)', result.stdout)
                if temps:
                    # Take average of all core temps
                    temp_c = sum(float(t) for t in temps) / len(temps)
            except (subprocess.TimeoutExpired, ValueError):
                pass

        return CPUMetrics(
            load_1m=load_1m,
            load_5m=load_5m,
            load_15m=load_15m,
            temp_c=temp_c,
            core_count=self._cpu_count
        )

    def read_mem(self) -> MemoryMetrics:
        """
        Read memory metrics from /proc/meminfo.
        """
        meminfo = {}
        try:
            with open("/proc/meminfo", "r") as f:
                for line in f:
                    parts = line.split()
                    if len(parts) >= 2:
                        key = parts[0].rstrip(":")
                        # Value in kB
                        meminfo[key] = int(parts[1])
        except (IOError, ValueError):
            pass

        total_kb = meminfo.get("MemTotal", 0)
        available_kb = meminfo.get("MemAvailable", 0)
        swap_total_kb = meminfo.get("SwapTotal", 0)
        swap_free_kb = meminfo.get("SwapFree", 0)

        total_mb = total_kb / 1024
        available_mb = available_kb / 1024
        used_mb = total_mb - available_mb
        swap_total_mb = swap_total_kb / 1024
        swap_used_mb = (swap_total_kb - swap_free_kb) / 1024

        pressure_pct = (used_mb / total_mb * 100) if total_mb > 0 else 0.0

        return MemoryMetrics(
            total_mb=total_mb,
            available_mb=available_mb,
            used_mb=used_mb,
            swap_total_mb=swap_total_mb,
            swap_used_mb=swap_used_mb,
            pressure_pct=pressure_pct
        )

    def read_io(self) -> IOMetrics:
        """
        Read IO metrics: IO wait from /proc/stat.
        """
        io_wait_pct = 0.0

        try:
            with open("/proc/stat", "r") as f:
                for line in f:
                    if line.startswith("cpu "):
                        # cpu user nice system idle iowait irq softirq steal guest guest_nice
                        parts = line.split()
                        if len(parts) >= 6:
                            user = int(parts[1])
                            nice = int(parts[2])
                            system = int(parts[3])
                            idle = int(parts[4])
                            iowait = int(parts[5])

                            total = user + nice + system + idle + iowait
                            if total > 0:
                                io_wait_pct = (iowait / total) * 100
                        break
        except (IOError, ValueError, IndexError):
            pass

        # Disk utilization would need iostat, keeping simple for now
        disk_util_pct = 0.0

        return IOMetrics(
            io_wait_pct=io_wait_pct,
            disk_util_pct=disk_util_pct
        )

    def snapshot(self) -> TelemetrySnapshot:
        """
        Take a complete telemetry snapshot.

        This is the primary interface - call this to get all metrics at once.
        """
        import time

        return TelemetrySnapshot(
            gpu=self.read_gpu(),
            cpu=self.read_cpu(),
            memory=self.read_mem(),
            io=self.read_io(),
            timestamp=time.time()
        )

    def info(self) -> Dict:
        """Return observer capabilities."""
        return {
            "nvidia_smi_available": self._has_nvidia_smi,
            "lm_sensors_available": self._has_lm_sensors,
            "cpu_count": self._cpu_count,
        }


# Module-level singleton
telemetry = TelemetryObserver()


def read_telemetry() -> TelemetrySnapshot:
    """Convenience function to get a telemetry snapshot."""
    return telemetry.snapshot()


if __name__ == "__main__":
    # Quick test
    import json
    obs = TelemetryObserver()
    print("Observer capabilities:", obs.info())
    snap = obs.snapshot()
    print("Telemetry snapshot:")
    print(json.dumps(snap.to_dict(), indent=2))
