package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"git.sophuwu.com/amdgpu-qstat/config"
)

func fatal(msg string) {
	if config.ShowErr {
		fmt.Fprintf(os.Stderr, "fatal error: %s\n", msg)
	}
	os.Exit(1)
}

var HwmonPath string

func getHwmonPath() {
	if HwmonPath != "" {
		return
	}
	hwmonPath := "/sys/class/hwmon"
	dr, err := os.ReadDir(hwmonPath)
	if err != nil || len(dr) == 0 {
		fatal("unable to read hwmon")
	}
	var b []byte
	for _, entry := range dr {
		b, err = os.ReadFile(filepath.Join(hwmonPath, entry.Name(), "name"))
		if err != nil {
			continue
		}
		if strings.TrimSpace(string(b)) == "amdgpu" {
			HwmonPath = filepath.Join(hwmonPath, entry.Name())
			return
		}
	}
	fatal("amdgpu hwmon not found")
}

func ERR(msg string) {
	if config.ShowErr {
		fmt.Fprintf(os.Stderr, "error: %s\n", msg)
	}
}

func main() {
	getHwmonPath()

	var f Freq
	if config.Freq || config.MemFreq {
		err := f.Get()
		if err != nil {
			ERR("unable to read frequency")
			config.Freq = false
			config.MemFreq = false
		}
	}

	var m Mem
	if config.Mem || config.MemPerc {
		err := m.Get()
		if err != nil {
			ERR("unable to read memory usage")
			config.Mem = false
			config.MemPerc = false
		}
	}

	var val int
	var err error
	var s string

	if config.Use || config.Temp || config.Freq || config.Power {
		s += "GPU: "
		if config.Use {
			val, err = readNum("device/gpu_busy_percent")
			if err != nil {
				ERR("unable to read GPU usage")
			} else {
				s += fmt.Sprintf("%3d%%  ", val)
			}
		}
		if config.Freq {
			if f.Core != "" {
				s += fmt.Sprintf("%s  ", f.Core)
			}
		}
		if config.Power {
			val, err = readNum("power1_input")
			if err != nil {
				ERR("unable to read power consumption")
			} else {
				s += fmt.Sprintf("%3.0f W  ", float64(val)/1000000.0)
			}
		}
		if config.Temp {
			val, err = readNum("temp1_input")
			if err != nil {
				ERR("unable to read temperature")
			} else {
				s += fmt.Sprintf("%3d C", val/1000)
			}
		}
		if config.Fan || config.MemFreq || config.Mem || config.MemPerc {
			s += "  | "
		}
	}
	if config.Mem || config.MemPerc || config.MemFreq {
		s += "MEM: "
		if config.Mem && config.MemPerc {
			s += fmt.Sprintf("%s (%s)  ", m.Gb, m.Perc)
		} else if config.Mem {
			s += fmt.Sprintf("%s  ", m.Gb)
		} else if config.MemPerc {
			s += fmt.Sprintf("%s  ", m.Perc)
		}
		if config.MemFreq {
			s += fmt.Sprintf("%s", f.Mem)
		}
		if config.Fan {
			s += "  | "
		}
	}
	if config.Fan {
		val, err = readNum("fan1_input")
		if err != nil {
			ERR("unable to read fan speed")
		} else {
			s += fmt.Sprintf("FAN: %4d RPM", val)
		}
	}

	fmt.Println(s)

}

type Mem struct {
	Perc string
	Gb   string
}

func (m *Mem) Get() error {
	used, err := readNum("device/mem_info_vram_used")
	if err != nil {
		return fmt.Errorf("unable to get memory usage")
	}
	total, err := readNum("device/mem_info_vram_total")
	if err != nil {
		return fmt.Errorf("unable to get memory usage")
	}
	Perc := float64(used) / float64(total) * 100
	m.Perc = fmt.Sprintf("%3.0f%%", Perc)
	m.Gb = fmt.Sprintf("%4.2f GB / %4.2f GB", float64(used)/1073741824.0, float64(total)/1073741824.0)
	return nil
}

func readNum(path string) (int, error) {
	b, err := os.ReadFile(filepath.Join(HwmonPath, path))
	if err != nil {
		return 0, err
	}
	var val int
	for _, c := range b {
		if c < '0' || c > '9' {
			break
		}
		val = val*10 + int(c-'0')
	}
	return int(val), nil
}

type Freq struct {
	Core string
	Mem  string
}

func (f *Freq) Get() error {
	dr, err := os.ReadDir(filepath.Join(HwmonPath))
	if err != nil {
		return fmt.Errorf("unable to get frequency")
	}
	var b []byte
	var val int
	var s string
	for _, entry := range dr {
		if strings.HasPrefix(entry.Name(), "freq") && strings.HasSuffix(entry.Name(), "_label") {
			b, err = os.ReadFile(filepath.Join(HwmonPath, entry.Name()))
			if err != nil {
				return fmt.Errorf("unable to get frequency")
			}
			val, err = readNum(strings.Replace(entry.Name(), "_label", "_input", 1))
			if err != nil {
				return fmt.Errorf("unable to get frequency")
			}
			s = fmt.Sprintf("%4.0f MHz", float64(val)/1000000.0)
			if strings.TrimSpace(string(b)) == "mclk" {
				f.Mem = s
			} else if strings.TrimSpace(string(b)) == "sclk" {
				f.Core = s
			}
		}
	}
	return nil
}