package main

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"testing"
	"time"
)

// fakeKFD builds the sysfs shape the workstation actually has: one directory
// per ROCm process, each holding a vram_<node> file. Sampled from the live box
// on 02-08-2026, where the CPT run appeared as proc/478104/vram_35881.
func fakeKFD(t *testing.T, vramByPID map[int]int64) string {
	t.Helper()
	root := t.TempDir()
	for pid, vram := range vramByPID {
		dir := filepath.Join(root, strconv.Itoa(pid))
		if err := os.MkdirAll(dir, 0o755); err != nil {
			t.Fatal(err)
		}
		f := filepath.Join(dir, "vram_35881")
		if err := os.WriteFile(f, []byte(strconv.FormatInt(vram, 10)+"\n"), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	return root
}

func TestForeignExcludesOurChild(t *testing.T) {
	root := fakeKFD(t, map[int]int64{478104: 12791693312, 999: 4096})
	p := probe{kfdRoot: root}

	all := p.foreign(0)
	if len(all) != 2 {
		t.Fatalf("with no child running, both processes are foreign, got %d", len(all))
	}

	ours := p.foreign(999)
	if len(ours) != 1 || ours[0].PID != 478104 {
		t.Fatalf("our own llama-server must not count as a contender, got %+v", ours)
	}
	if ours[0].VRAM != 12791693312 {
		t.Errorf("per-process VRAM = %d, want the value from vram_35881", ours[0].VRAM)
	}
}

// The transcriber is a ROCm process on the same card, so it registers on the
// KFD exactly like a contender does. Reading it as one is what happened on
// 2026-08-09 while CW2 ran under its own systemd unit: mavgpud yielded, waited
// five polls, loaded the model, yielded again, and never held it for a whole
// minute. Excluding every child is the fix and this is the test of it.
func TestForeignExcludesEveryChild(t *testing.T) {
	p := probe{kfdRoot: fakeKFD(t, map[int]int64{478104: 12791693312, 999: 4096, 1001: 1717986918})}

	ours := p.foreign(999, 1001)
	if len(ours) != 1 || ours[0].PID != 478104 {
		t.Fatalf("only the CPT run is a contender, got %+v", ours)
	}
	// A child that is not running reports pid 0, which must exclude nothing.
	if got := p.foreign(999, 0); len(got) != 2 {
		t.Errorf("a stopped child excludes nobody: got %d contenders, want 2", len(got))
	}
}

// An empty KFD tree is the state that permits a start, so it must read as empty
// rather than as an error the caller has to interpret.
func TestForeignEmptyAndMissing(t *testing.T) {
	if got := (probe{kfdRoot: t.TempDir()}).foreign(0); len(got) != 0 {
		t.Errorf("empty kfd tree: got %d processes, want 0", len(got))
	}
	if got := (probe{kfdRoot: "/nonexistent"}).foreign(0); got != nil {
		t.Errorf("missing kfd tree: got %+v, want nil", got)
	}
}

func TestFreeVRAM(t *testing.T) {
	dev := t.TempDir()
	write := func(name, v string) {
		if err := os.WriteFile(filepath.Join(dev, name), []byte(v), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	// The live numbers from the workstation while the CPT run held the card.
	write("mem_info_vram_total", "17163091968\n")
	write("mem_info_vram_used", "13396389888\n")
	p := probe{drmDev: dev}
	if got, want := p.freeVRAM(), int64(3766702080); got != want {
		t.Errorf("freeVRAM = %d, want %d", got, want)
	}
	if got := (probe{drmDev: "/nonexistent"}).freeVRAM(); got != 0 {
		t.Errorf("unreadable card reports %d free, want 0 so nothing starts", got)
	}
}

// With no model loaded the supervisor must still answer, and it must answer 503
// rather than hanging or proxying into a closed port. Maven reads this endpoint
// on a timer forever, including while the workstation is busy.
func TestHealthAndProxyRefuseWhenNotReady(t *testing.T) {
	s := &supervisor{run: newRunner("fake", "/bin/true", nil, "")}
	h := s.handler(mustURL(t, "http://127.0.0.1:1"))

	for _, path := range []string{"/health", "/v1/chat/completions"} {
		w := httptest.NewRecorder()
		h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil))
		if w.Code != http.StatusServiceUnavailable {
			t.Errorf("%s with no model: got %d, want 503", path, w.Code)
		}
	}
}

func mustURL(t *testing.T, s string) *url.URL {
	t.Helper()
	u, err := url.Parse(s)
	if err != nil {
		t.Fatal(err)
	}
	return u
}

// Yielding is all or nothing. A CPT run wants the whole card, so handing back
// the language model while the transcriber keeps 1.6GB mapped would leave the
// other job failing its allocation, which is the outcome yielding exists to
// prevent.
func TestYieldStopsEveryChild(t *testing.T) {
	idle := "while : ; do sleep 1 ; done"
	s := &supervisor{
		cfg:   config{EvictAfter: 1, StopGrace: duration(2 * time.Second)},
		probe: probe{kfdRoot: fakeKFD(t, map[int]int64{478104: 12791693312})},
		run:   newRunner("llama-server", fakeServer(t, idle), nil, ""),
		stt:   newRunner("cw2", fakeServer(t, idle), nil, ""),
	}
	for _, r := range s.children() {
		if err := r.start(); err != nil {
			t.Fatal(err)
		}
	}
	s.tick(context.Background())
	for _, r := range s.children() {
		if r.running() {
			t.Errorf("%s outlived the yield", r.name)
		}
	}
}