package main

import (
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// The card is an AMD 7900 GRE with 16GB, driven by amdgpu and ROCm. Everything
// here reads sysfs and forks nothing: rocm-smi is not even installed on the
// workstation, and a poll that costs a subprocess every second is a poll that
// gets tuned down until it is useless.

// gpuProc — one process holding the compute engine.
type gpuProc struct {
	PID  int
	Comm string
	VRAM int64 // bytes, as the kernel accounts them to this process
}

// probe reads the two sysfs trees the supervisor decides from.
//
// kfdRoot is /sys/class/kfd/kfd/proc, one directory per ROCm process. The
// directory appears when the process initialises HIP, which is well before it
// allocates anything large. That is the whole reason this works: the job that
// is about to want the card announces itself while it is still starting up,
// so we see the contender rather than only the winner of an allocation race.
//
// drmDev is /sys/class/drm/cardN/device, which reports total and used VRAM for
// the card as a whole.
type probe struct {
	kfdRoot string
	drmDev  string
}

// foreign lists every ROCm process that is not ours. selfPID is the supervisor's
// llama-server child, or 0 when it is not running.
//
// An unreadable kfd tree returns no processes and no error. That is deliberate
// and it is the safe direction only because startVRAM also has to agree before
// anything launches: a supervisor that cannot see the KFD never sees free VRAM
// either, because the CPT run holding the card shows up in the drm totals.
func (p probe) foreign(selfPID int) []gpuProc {
	entries, err := os.ReadDir(p.kfdRoot)
	if err != nil {
		return nil
	}
	var out []gpuProc
	for _, e := range entries {
		pid, err := strconv.Atoi(e.Name())
		if err != nil || pid == selfPID {
			continue
		}
		out = append(out, gpuProc{
			PID:  pid,
			Comm: readComm(pid),
			VRAM: p.procVRAM(e.Name()),
		})
	}
	return out
}

// procVRAM sums the per-node vram_* files under one process directory. The
// suffix is the KFD topology node id (vram_35881 on this card), so it is
// globbed rather than named, and a machine with two cards sums both.
func (p probe) procVRAM(pid string) int64 {
	matches, err := filepath.Glob(filepath.Join(p.kfdRoot, pid, "vram_*"))
	if err != nil {
		return 0
	}
	var total int64
	for _, m := range matches {
		total += readInt(m)
	}
	return total
}

// freeVRAM reports the bytes the card has left. Used only to decide whether to
// start: a shortfall here means llama-server would refuse to load anyway. It is
// never used to decide to stop, because by the time free VRAM has dropped the
// other job has already failed its allocation, which is exactly the outcome
// yielding exists to prevent.
func (p probe) freeVRAM() int64 {
	total := readInt(filepath.Join(p.drmDev, "mem_info_vram_total"))
	used := readInt(filepath.Join(p.drmDev, "mem_info_vram_used"))
	if total <= 0 {
		return 0
	}
	if free := total - used; free > 0 {
		return free
	}
	return 0
}

func readInt(path string) int64 {
	b, err := os.ReadFile(path)
	if err != nil {
		return 0
	}
	n, err := strconv.ParseInt(strings.TrimSpace(string(b)), 10, 64)
	if err != nil {
		return 0
	}
	return n
}

// readComm names the contender for the log. The log is the instrument for the
// open question in Vikunja #488: whether a process can want this card without
// ever registering on the KFD, which a Vulkan or video-decode job would.
func readComm(pid int) string {
	b, err := os.ReadFile(filepath.Join("/proc", strconv.Itoa(pid), "comm"))
	if err != nil {
		return "?"
	}
	return strings.TrimSpace(string(b))
}