import pynvml
import time
from colorama import Fore, Style, init
import os

# Initialize colorama
init(autoreset=True)



def monitor_gpu_ram_usage(interval=2, threshold_gb=2):
    pynvml.nvmlInit()
    # Initialize NVML
    try:
        device_count = pynvml.nvmlDeviceGetCount()
        print(f"Found {device_count} GPU(s).")

        while True:
            for i in range(device_count):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                info = pynvml.nvmlDeviceGetMemoryInfo(handle)



                print(f"GPU {i}:")
                print(f"  Total RAM: {info.total / 1024 ** 2:.2f} MB")
                if(info.used / 1024 ** 2 >= 2.5 * 1024 ):
                    print(Fore.RED + f"  Used RAM: {info.used / 1024 ** 2:.2f} MB")
                    os.system("aplay /home/rog/repos/Tracker/NE-Smart-Tracker/Oxygen-Sys-Warning.wav")
                else:
                    print(f"  Used RAM: {info.used / 1024 ** 2:.2f} MB")
                print(f"  Free RAM: {info.free / 1024 ** 2:.2f} MB")
                print(Fore.GREEN + "-" * 30)
                print(Fore.GREEN)

            time.sleep(interval)  # Wait for the specified interval before checking again

    except KeyboardInterrupt:
        print("Monitoring stopped by user.")

    finally:
        # Shutdown NVML
        pynvml.nvmlShutdown()

if __name__ == "__main__":
    monitor_gpu_ram_usage(interval=2, threshold_gb=2)  # Check every 2 seconds, threshold is 2 GB