#!/usr/bin/env python3
"""
Real-time GPU and RAM monitoring utility
"""

import psutil
import time
import threading
from rich.console import Console
from rich.live import Live
from rich.table import Table

try:
    import pynvml
    NVIDIA_AVAILABLE = True
except ImportError:
    NVIDIA_AVAILABLE = False

console = Console()

class SystemMonitor:
    def __init__(self):
        self.running = False
        self.gpu_available = False
        
        if NVIDIA_AVAILABLE:
            try:
                pynvml.nvmlInit()
                self.gpu_count = pynvml.nvmlDeviceGetCount()
                self.gpu_available = self.gpu_count > 0
            except:
                pass
    
    def get_ram_usage(self):
        memory = psutil.virtual_memory()
        return {
            "used_gb": memory.used / 1024**3,
            "total_gb": memory.total / 1024**3,
            "percent": memory.percent
        }
    
    def get_gpu_usage(self):
        if not self.gpu_available:
            return None
        
        gpu_info = []
        try:
            for i in range(self.gpu_count):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                
                gpu_info.append({
                    "id": i,
                    "used_gb": info.used / 1024**3,
                    "total_gb": info.total / 1024**3,
                    "percent": (info.used / info.total) * 100
                })
        except:
            pass
        
        return gpu_info
    
    def create_table(self):
        table = Table(title="System Resources")
        table.add_column("Resource", style="cyan")
        table.add_column("Used", style="yellow")
        table.add_column("Total", style="green")
        table.add_column("Usage %", style="red")
        
        # RAM
        ram = self.get_ram_usage()
        table.add_row(
            "RAM",
            f"{ram['used_gb']:.2f} GB",
            f"{ram['total_gb']:.2f} GB",
            f"{ram['percent']:.1f}%"
        )
        
        # GPU
        if self.gpu_available:
            gpus = self.get_gpu_usage()
            if gpus:
                for gpu in gpus:
                    table.add_row(
                        f"GPU {gpu['id']}",
                        f"{gpu['used_gb']:.2f} GB",
                        f"{gpu['total_gb']:.2f} GB",
                        f"{gpu['percent']:.1f}%"
                    )
        else:
            table.add_row("GPU", "N/A", "N/A", "N/A")
        
        return table
    
    def start_monitoring(self):
        self.running = True
        
        with Live(self.create_table(), refresh_per_second=1, console=console) as live:
            while self.running:
                live.update(self.create_table())
                time.sleep(1)
    
    def stop_monitoring(self):
        self.running = False
        if self.gpu_available:
            try:
                pynvml.nvmlShutdown()
            except:
                pass

def monitor_resources():
    """Start monitoring in a separate thread"""
    monitor = SystemMonitor()
    thread = threading.Thread(target=monitor.start_monitoring, daemon=True)
    thread.start()
    return monitor

if __name__ == "__main__":
    console.print("[bold cyan]System Resource Monitor[/bold cyan]")
    console.print("Press Ctrl+C to stop\n")
    
    monitor = SystemMonitor()
    try:
        monitor.start_monitoring()
    except KeyboardInterrupt:
        console.print("\n[yellow]Monitoring stopped.[/yellow]")
        monitor.stop_monitoring()