import subprocess
import re
import os

class HardwareMonitor:
    """
    Tracks GPU health and temperature to prevent hardware damage.
    Uses nvidia-smi as a zero-dependency fallback.
    """
    def __init__(self):
        self.has_gpu = torch_has_gpu = False
        try:
            import torch
            self.has_gpu = torch.cuda.is_available()
        except:
            pass

    def get_gpu_stats(self):
        """Returns dict of (temp, memory_used, power_draw)."""
        if not self.has_gpu:
            return {"temp": 0, "mem": 0, "power": 0, "error": "No GPU"}

        try:
            # Query nvidia-smi for specific metrics
            cmd = ["nvidia-smi", "--query-gpu=temperature.gpu,memory.used,power.draw", "--format=csv,noheader,nounits"]
            result = subprocess.check_output(cmd, encoding='utf-8')
            temp, mem, power = result.strip().split(',')
            
            return {
                "temp": int(temp),
                "mem": int(mem),
                "power": float(power),
                "error": None
            }
        except Exception as e:
            return {"temp": 0, "mem": 0, "power": 0, "error": str(e)}

    def check_safety(self, max_temp=80):
        """Returns (is_safe, message)."""
        stats = self.get_gpu_stats()
        if stats["error"]:
             return True, "No GPU monitoring available."
             
        if stats["temp"] >= max_temp + 5:
            return False, f"CRITICAL: GPU Temp {stats['temp']}°C exceeds shutdown limit!"
        elif stats["temp"] >= max_temp:
            return False, f"WARNING: GPU Temp {stats['temp']}°C. Throttling/Pausing recommended."
            
        return True, "Hardware health normal."

if __name__ == "__main__":
    monitor = HardwareMonitor()
    print(monitor.get_gpu_stats())
    print(monitor.check_safety())