package main

import (
	"strconv"
	"strings"
	"sync"
	"time"
)

// Memory accounting must match how the machine is built:
//
//   - Unified memory (DGX Spark / GB10 "PGX", Grace-Hopper style): the GPU uses the same memory as the CPU.
//     nvidia-smi reports memory as [N/A], and the right budget is the system RAM (/proc/meminfo).
//   - Discrete GPUs (most DGX servers and workstations): models live in each GPU's own memory (VRAM). The
//     system RAM can be terabytes while a model still does not fit, so the budget is the GPUs' free VRAM.

type gpuMem struct {
	Total, Used uint64 // summed over all GPUs
	MaxOne      uint64 // total memory of the largest single GPU (vLLM's fraction is per GPU)
	GPUs        int
}

func (g gpuMem) avail() uint64 {
	if g.Used > g.Total {
		return 0
	}
	return g.Total - g.Used
}

// parseGPUMem reads `nvidia-smi --query-gpu=memory.total,memory.used --format=csv,noheader,nounits`.
// ok is false when no GPU reports numbers (no GPU, or unified memory showing [N/A]).
func parseGPUMem(out string) (g gpuMem, ok bool) {
	for _, line := range strings.Split(strings.TrimSpace(out), "\n") {
		f := strings.Split(line, ",")
		if len(f) < 2 {
			continue
		}
		tot, err1 := strconv.ParseUint(strings.TrimSpace(f[0]), 10, 64)
		used, err2 := strconv.ParseUint(strings.TrimSpace(f[1]), 10, 64)
		if err1 != nil || err2 != nil { // "[N/A]" or "[Not Supported]"
			continue
		}
		tot <<= 20 // MiB -> bytes
		used <<= 20
		g.Total += tot
		g.Used += used
		g.GPUs++
		if tot > g.MaxOne {
			g.MaxOne = tot
		}
	}
	return g, g.GPUs > 0
}

var gpuMemCache = struct {
	sync.Mutex
	at time.Time
	g  gpuMem
	ok bool
}{}

// readGPUMem asks nvidia-smi (cached for a few seconds: it is called on every admission decision).
func readGPUMem() (gpuMem, bool) {
	gpuMemCache.Lock()
	defer gpuMemCache.Unlock()
	if time.Since(gpuMemCache.at) < 3*time.Second {
		return gpuMemCache.g, gpuMemCache.ok
	}
	out, err := runCmd(3*time.Second, "nvidia-smi", "--query-gpu=memory.total,memory.used", "--format=csv,noheader,nounits")
	gpuMemCache.at = time.Now()
	if err != nil {
		gpuMemCache.g, gpuMemCache.ok = gpuMem{}, false
	} else {
		gpuMemCache.g, gpuMemCache.ok = parseGPUMem(out)
	}
	return gpuMemCache.g, gpuMemCache.ok
}

// budgetMemory is the memory models compete for: GPU memory on discrete-GPU machines,
// system RAM on unified-memory machines (and CPU-only ones). perGPU is the largest single GPU's
// size, or 0 when the budget is system RAM.
func budgetMemory() (total, avail, perGPU uint64) {
	if g, ok := readGPUMem(); ok {
		return g.Total, g.avail(), g.MaxOne
	}
	t, a := memInfo()
	return t, a, 0
}