import GPUtil
import matplotlib.pyplot as plt
import os
import psutil
import gradio as gr


# TODO: Use plotly trace to update a plot of the stats; try to get near parity to system monitor
# Ask: what stats are users most interested in knowing?
def get_stats(plot):
    stats = ""

    # Get CPU usage
    cpu_percent = psutil.cpu_percent(interval=0)

    # Get the number of CPU cores
    num_cores = os.cpu_count()

    # Get RAM usage
    ram = psutil.virtual_memory()
    total_ram = psutil.virtual_memory().total

    stats += f"CPU Usage: {cpu_percent}%\n"
    stats += f"Number of CPU Cores: {num_cores}\n"
    stats += f"RAM Usage: {ram.percent}%\n"
    stats += f"RAM Used: {ram.used / (1024 ** 3):.2f} GB\n"
    stats += f"Total RAM: {total_ram / (1024 ** 3):.2f} GB\n"
    # print(f"CPU Usage: {cpu_percent}%")
    # print(f"Number of CPU Cores: {num_cores}")
    # print(f"RAM Usage: {ram.percent}%")
    # print(f"RAM Used: {ram.used / (1024 ** 3):.2f} GB")
    # print(f"Total RAM: {total_ram / (1024 ** 3):.2f} GB")

    # Get GPU usage and memory info
    gpu_info = GPUtil.getGPUs()

    for i, gpu in enumerate(gpu_info):
        stats += f"GPU {i + 1}:\n"
        stats += f"  GPU Name: {gpu.name}\n"
        stats += f"  GPU Usage: {gpu.load * 100}%\n"
        stats += f"  GPU VRAM Usage: {(gpu.memoryUtil * 100):.2f}%\n"
        stats += f"  GPU Free VRAM: {(gpu.memoryFree / 1024):.2f} MB\n"
        stats += f"  GPU Total VRAM: {gpu.memoryTotal / 1024} MB\n\n"
        # print(f"GPU {i + 1}:")
        # print(f"  GPU Name: {gpu.name}")
        # print(f"  GPU Usage: {gpu.load * 100}%")
        # print(f"  GPU VRAM Usage: {(gpu.memoryUtil * 100):.2f}%")
        # print(f"  GPU Free VRAM: {(gpu.memoryFree / 1024):.2f} MB")
        # print(f"  GPU Total VRAM: {gpu.memoryTotal / 1024} MB\n")

    return stats, None


with gr.Blocks() as demo:
    gr.Markdown(value='# System Monitor')
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown('## View system resource usage for the demo')
            stats_box = gr.TextArea(label="System Stats")
        with gr.Column(scale=1):
            plot = gr.Plot(label="Plot", value=None)

    dep = demo.load(get_stats, [plot], [stats_box, plot], every=5)
    stats_box.change(get_stats, [plot], [stats_box, plot], every=5, cancels=[dep])

    with gr.Accordion('Developer Notes:', open=False):
        gr.Markdown('I\'ve seen a number of Ai tools that report system resource usage and think that is a good trend.\n\n'
                    'With more time I could see a "Task Manager" component for gradio being useful for demos.\n\n'
                    'The plot for this demo is broken but I may revisit this later.')


if __name__ == '__main__':
    demo.queue().launch()