package stat

// Modified from https://github.com/influxdata/telegraf/blob/master/plugins/inputs/amd_rocm_smi/amd_rocm_smi.go
// Original License: MIT

import (
	"encoding/json"
	"errors"
	"os"
	"os/exec"
	"strconv"
)

type ROCmSMI struct {
	BinPath string
}

func (rsmi *ROCmSMI) Gather() ([]float64, error) {
	data := rsmi.pollROCmSMI()

	return gatherROCmSMI(data)
}

func (rsmi *ROCmSMI) Start() error {
	if _, err := os.Stat(rsmi.BinPath); os.IsNotExist(err) {
		binPath, err := exec.LookPath("rocm-smi")
		if err != nil {
			return errors.New("Didn't find the adequate tool to query GPU utilization")
		}
		rsmi.BinPath = binPath
	}
	return nil
}

func (rsmi *ROCmSMI) pollROCmSMI() []byte {
	cmd := exec.Command(rsmi.BinPath,
		"-u",
		"--json",
	)
	gs, err := cmd.CombinedOutput()
	if err != nil {
		return nil
	}
	return gs
}

func gatherROCmSMI(ret []byte) ([]float64, error) {
	var gpus map[string]GPU
	var percentage []float64

	err := json.Unmarshal(ret, &gpus)
	if err != nil {
		return nil, err
	}

	for _, gpu := range gpus {
		gp, _ := strconv.ParseFloat(gpu.GpuUsePercentage, 64)
		percentage = append(percentage, gp)
	}

	return percentage, nil
}

type GPU struct {
	GpuUsePercentage string `json:"GPU use (%)"`
}