#!/usr/bin/env python3

from time import sleep
import json
import psutil
import subprocess
import os
from dotenv import dotenv_values

amdgpuinfo_installed = False
# Try to import GPU info
try:
    import pyamdgpuinfo as gpuinfo
    amdgpuinfo_installed = True
except:
    pass

def cpu():
    util = psutil.cpu_percent(percpu=False)
    per_core = psutil.cpu_percent(percpu=True)
    return {
        "avg_display": f"{int(util):2}",
        "avg": util,
        "cores": [{"value": core, "display": f"{int(core):2}"} for core in per_core],
        "cores_display": "".join([f"{int(core):3}" for core in per_core]),
        "core_count": len(per_core),
    }

def gpu():
    if not amdgpuinfo_installed or gpuinfo.detect_gpus() <= 0:
        return {
            "available": False,
            "load": 0,
            "power_state": 0,
            "memory": 0
        }
    gpu = gpuinfo.get_gpu(0)
    return {
        "available": True,
        "load": gpu.query_load(),
        "power_state": gpu.query_power(),
        "memory": gpu.query_vram_usage() / gpu.memory_info['vram_size']
    }


def sensor(chip: str):
    return {
        temp.label: temp.current for temp in psutil.sensors_temperatures().get(chip, [])
    }


def memory():
    mem = psutil.virtual_memory()
    return {
        "total": mem.total,
        "available": mem.available,
        "percent": mem.percent,
        "used": mem.used,
        "free": mem.free,
        "active": mem.active,
        "inactive": mem.inactive,
        "buffers": mem.buffers,
        "cached": mem.cached,
        "shared": mem.shared,
        "slab": mem.slab,
    }


def swap():
    swap = psutil.swap_memory()
    return {
        "total": swap.total,
        "used": swap.used,
        "free": swap.free,
        "percent": swap.percent,
    }


def reboot():
    running = os.uname().release
    version_proc = subprocess.run(["pacman", "-Q", "linux"], capture_output=True)
    installed = version_proc.stdout.strip().decode("utf-8").split(" ")[1]
    return running != installed


sensor_list = ["coretemp"]

counter = 60
machine_info = dotenv_values("/etc/machine-info")

while True:
    # Refresh machine info every minute, in case it changes
    if counter == 0:
        machine_info = dotenv_values("/etc/machine-info")
        counter = 60

    result = {
        "cpu": cpu(),
        "gpu": gpu(),
        "sensors": {chip: sensor(chip) for chip in sensor_list},
        "memory": memory(),
        "swap": swap(),
        "reboot": reboot(),
        "hostname": machine_info.get("PRETTY_HOSTNAME"),
    }
    print(json.dumps(result), flush=True)
    counter -= 1
    sleep(1)