#!/usr/bin/python3

import subprocess
import psutil
import re

class Load:
    def __init__(self):
        self.nvidia = False
        self.amd = False
        try:
            gpu_info = subprocess.Popen(["nvidia-smi", "-q", "-d", "temperature"], stdout=subprocess.PIPE, universal_newlines=True)
        except:
            pass
        else:
            self.nvidia = True
        if self.nvidia == False:
            try:
                gpu_info = subprocess.Popen(["sensors", "amdgpu-*"], stdout=subprocess.PIPE, universal_newlines=True)
            except:
                pass
            else:
                self.amd = True


        if self.nvidia == True:
            self.update_nvidia()
        elif self.amd == True:
            self.update_amd()
        else:
            print('wtf')

        self.update_cpu()

    def update_gpu(self):
        if self.nvidia == True:
            self.update_nvidia()
        elif self.amd == True:
            self.update_amd()       

    def update_nvidia(self):
        #gpu_info = subprocess.Popen(["nvidia-smi", "-q", "-d", "temperature"], stdout=subprocess.PIPE, universal_newlines=True)
        gpu_info = subprocess.Popen(["nvidia-smi", "-q"], stdout=subprocess.PIPE, universal_newlines=True)
        for info in iter(gpu_info.stdout.readline, ''):
            if 'GPU Current Temp' in info:
                info = re.findall('[0-9]+', info)
                info = info[0]
                self.current_temp = int(info)
            elif 'Gpu' in info:
                info = re.findall('[0-9]+', info)
                info = info[0]
                self.utilization = int(info)
            #    info = info.replace('\n', '')
            #    info = re.findall('[0-9]+', info)
            #    info = info[0]
            #    self.current_temp = info
            #elif "Slowdown Temp" in info:
            #    info = info.replace('\n', '')
            #    info = re.findall('[0-9]+', info)
            #    info = info[0]
            #    self.slowdown_temp = info
            else:
                pass
    
    def update_amd(self):
        gpu_info = subprocess.Popen(["sensors", "amdgpu-*"], stdout=subprocess.PIPE, universal_newlines=True)
        for info in iter(gpu_info.stdout.readline, ''):
            if 'edge' in info or 'temp1' in info:
                info = info.replace('\n', '')
                info = re.findall(r'\+\d*', info)
                info[0] = info[0].replace('+', '')
                info[1] = info[1].replace('+', '')
                self.current_temp = int(info[0])
                self.slowdown_temp = int(info[1])
                self.dayum = subprocess.run('cat /sys/class/drm/card0/device/gpu_busy_percent', shell=True, capture_output=True)
                self.dayum = re.findall('[0-9]+', str(self.dayum.stdout))
                self.utilization = self.dayum[0]
        #cat /sys/class/drm/card0/device/gpu_busy_percent
        #cat /sys/kernel/debug/dri/0/amdgpu_pm_info

    def update_cpu(self):
        dick = psutil.sensors_temperatures()
        for sensor_list in dick.values():
            self.processor_temp = int(sensor_list[0][1])
        

