import psutil
import time
import GPUtil
from rich.live import Live
from rich.progress import Progress, BarColumn, ProgressColumn
from rich.console import Console
from rich.text import Text

class PercentageColumn(ProgressColumn):
    def render(self, task):
        return Text(f"{task.percentage:>5.1f}%")

def main():
    console = Console()
    progress = Progress(
        "[progress.description]{task.description}",
        BarColumn(bar_width=None),
        PercentageColumn(),
    )

    # Create tasks for CPU and GPUs
    cpu_task = progress.add_task("CPU Utilization", total=100)
    gpu_tasks = []
    gpus = GPUtil.getGPUs()
    for gpu in gpus:
        gpu_task = progress.add_task(f"GPU {gpu.id} Utilization", total=100)
        vram_task = progress.add_task(f"GPU {gpu.id} VRAM Usage", total=100)
        gpu_tasks.append((gpu.id, gpu_task, vram_task))

    with Live(progress, refresh_per_second=10):
        try:
            while True:
                # Update CPU utilization
                cpu_utilization = psutil.cpu_percent(interval=None)
                progress.update(cpu_task, completed=cpu_utilization)

                # Update GPU info
                gpus = GPUtil.getGPUs()
                for gpu in gpus:
                    gpu_util = gpu.load * 100
                    mem_util = gpu.memoryUtil * 100
                    for gid, gpu_task_id, vram_task_id in gpu_tasks:
                        if gid == gpu.id:
                            progress.update(gpu_task_id, completed=gpu_util)
                            progress.update(vram_task_id, completed=mem_util)

                time.sleep(0.1)
        except KeyboardInterrupt:
            pass

if __name__ == '__main__':
    main()