package main

import (
	"strings"
	"testing"
	"time"
)

var t0 = time.Unix(10000, 0)

func at(sec int) time.Time { return t0.Add(time.Duration(sec) * time.Second) }

func newG() *ownerGuard {
	return newOwnerGuard(guardConfig{Enabled: true, ActiveIdleS: 30, ReturnIdleS: 300, GraceS: 30})
}

func TestAvailableWhenOwnerIsAway(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 3600})
	if d.Mode != modeAvailable || d.StopAccepting {
		t.Fatalf("%+v", d)
	}
}

func TestUnknownIdleIsNotAnOwnerSignal(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: -1})
	if d.Mode != modeAvailable {
		t.Fatalf("headless machine must stay available: %+v", d)
	}
}

func TestOwnerActivityDrainsThenUnloads(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 5, InFlight: 2})
	if d.Mode != modeDraining || !d.StopAccepting || d.UnloadAll || d.CancelAll {
		t.Fatalf("should drain first, not unload yet: %+v", d)
	}
	d = g.step(guardInput{Now: at(10), OwnerIdleS: 5, InFlight: 2})
	if d.Mode != modeDraining || d.CancelAll {
		t.Fatalf("still within grace: %+v", d)
	}
	d = g.step(guardInput{Now: at(31), OwnerIdleS: 5, InFlight: 2})
	if d.Mode != modeOwner || !d.CancelAll || !d.UnloadAll {
		t.Fatalf("after the grace period running work is cancelled and models unloaded: %+v", d)
	}
}

func TestNoRunningWorkMeansImmediateUnload(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 5, InFlight: 0})
	if d.Mode != modeOwner || !d.UnloadAll || d.CancelAll || !d.StopAccepting {
		t.Fatalf("%+v", d)
	}
}

func TestRunningWorkFinishingEndsDrainEarly(t *testing.T) {
	g := newG()
	g.step(guardInput{Now: at(0), OwnerIdleS: 5, InFlight: 1})
	d := g.step(guardInput{Now: at(4), OwnerIdleS: 5, InFlight: 0})
	if d.Mode != modeOwner || d.CancelAll || !d.UnloadAll {
		t.Fatalf("%+v", d)
	}
}

func TestUnloadIsRequestedOnlyOncePerEpisode(t *testing.T) {
	g := newG()
	g.step(guardInput{Now: at(0), OwnerIdleS: 5})
	d := g.step(guardInput{Now: at(5), OwnerIdleS: 6})
	if d.UnloadAll {
		t.Fatalf("second step must not unload again: %+v", d)
	}
	if !d.StopAccepting || d.Mode != modeOwner {
		t.Fatalf("%+v", d)
	}
}

func TestComesBackOnlyAfterAQuietPeriod(t *testing.T) {
	g := newG()
	g.step(guardInput{Now: at(0), OwnerIdleS: 5})
	d := g.step(guardInput{Now: at(100), OwnerIdleS: 100}) // owner stopped 100 s ago: not enough
	if d.Mode != modeOwner || !d.StopAccepting {
		t.Fatalf("should still hold: %+v", d)
	}
	d = g.step(guardInput{Now: at(400), OwnerIdleS: 400})
	if d.Mode != modeAvailable || d.StopAccepting {
		t.Fatalf("quiet for 400 s -> available again: %+v", d)
	}
}

func TestMouseWiggleDoesNotFlapBetweenThresholds(t *testing.T) {
	g := newG()
	g.step(guardInput{Now: at(0), OwnerIdleS: 5})
	// idle 60 s: above "active" (30) but below "return" (300): must remain in owner mode
	d := g.step(guardInput{Now: at(60), OwnerIdleS: 60})
	if d.Mode != modeOwner {
		t.Fatalf("hysteresis broken: %+v", d)
	}
}

func TestForeignGpuWorkCountsAsOwner(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 3600, ForeignGPU: []string{"python3"}})
	if d.Mode != modeOwner || !d.UnloadAll {
		t.Fatalf("%+v", d)
	}
	// the process ends; it needs a quiet period of its own before we come back
	d = g.step(guardInput{Now: at(60), OwnerIdleS: 3600})
	if d.Mode != modeOwner {
		t.Fatalf("too early to return: %+v", d)
	}
	d = g.step(guardInput{Now: at(301), OwnerIdleS: 3600})
	if d.Mode != modeAvailable {
		t.Fatalf("should be back: %+v", d)
	}
}

func TestManualPauseWinsEvenWhenOwnerIsAway(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 9999, Paused: true})
	if d.Mode != modePaused || !d.StopAccepting || !d.UnloadAll {
		t.Fatalf("%+v", d)
	}
	d = g.step(guardInput{Now: at(5), OwnerIdleS: 9999})
	if d.Mode != modeAvailable {
		t.Fatalf("resume while the owner is away -> available: %+v", d)
	}
}

func TestManualPauseDrainsRunningWork(t *testing.T) {
	g := newG()
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 9999, Paused: true, InFlight: 3})
	if d.Mode != modeDraining || d.UnloadAll {
		t.Fatalf("%+v", d)
	}
	d = g.step(guardInput{Now: at(31), OwnerIdleS: 9999, Paused: true, InFlight: 3})
	if d.Mode != modePaused || !d.CancelAll || !d.UnloadAll {
		t.Fatalf("%+v", d)
	}
}

func TestResumeWhileOwnerIsActiveDoesNotMakeItAvailable(t *testing.T) {
	g := newG()
	g.step(guardInput{Now: at(0), OwnerIdleS: 5, Paused: true})
	d := g.step(guardInput{Now: at(5), OwnerIdleS: 5}) // resumed but the owner is still at the keyboard
	if d.Mode == modeAvailable || !d.StopAccepting {
		t.Fatalf("%+v", d)
	}
}

func TestDisabledGuardIgnoresActivityButNotManualPause(t *testing.T) {
	g := newOwnerGuard(guardConfig{Enabled: false, ActiveIdleS: 30, ReturnIdleS: 300, GraceS: 30})
	d := g.step(guardInput{Now: at(0), OwnerIdleS: 1, ForeignGPU: []string{"x"}})
	if d.Mode != modeAvailable {
		t.Fatalf("%+v", d)
	}
	d = g.step(guardInput{Now: at(1), OwnerIdleS: 1, Paused: true})
	if d.Mode != modePaused {
		t.Fatalf("manual pause must still work: %+v", d)
	}
}

func TestOtherGPUProcsIgnoresEveryRuntimeWeRun(t *testing.T) {
	procs := []gpuProc{
		{1, "ollama"}, {2, "/usr/local/bin/ollama"}, {3, "ollama_llama_server"},
		{4, "/home/pgx/.lmstudio/extensions/backends/llama.cpp-linux-arm64-nvidia-cuda13-2.41.0/llama-server"},
		{5, "VLLM::EngineCore"}, {6, "VLLM::Worker_TP0"},
		{7, "python3"}, {8, "/home/me/train.py"},
	}
	got := otherGPUProcs(procs, nil)
	if len(got) != 2 || got[0] != "python3" || got[1] != "train.py" {
		t.Fatalf("only the owner's own jobs should remain: %v", got)
	}
}

func TestVLLMEngineIsNotTheOwner(t *testing.T) {
	// regression: the guard once saw "VLLM::EngineCore" as foreign GPU work and unloaded the model being loaded
	if got := otherGPUProcs([]gpuProc{{10, "VLLM::EngineCore"}}, nil); len(got) != 0 {
		t.Fatalf("vLLM's own worker must be ignored: %v", got)
	}
}

func TestProcessesInOurServiceCgroupsAreIgnored(t *testing.T) {
	cg := func(pid int) string {
		if pid == 1 {
			return "0::/system.slice/ollama.service"
		}
		if pid == 2 {
			return "0::/user.slice/user-1000.slice/user@1000.service/app.slice/lmstudio.service"
		}
		return "0::/user.slice/user-1000.slice/session-2.scope"
	}
	got := otherGPUProcs([]gpuProc{{1, "mystery"}, {2, "node"}, {3, "blender"}}, cg)
	if len(got) != 1 || got[0] != "blender" {
		t.Fatalf("%v", got)
	}
}

func TestRetryCountdownFollowsTheQuietPeriod(t *testing.T) {
	g := newG()
	g.step(guardInput{Now: at(0), OwnerIdleS: 5})
	d := g.step(guardInput{Now: at(1), OwnerIdleS: 6})
	if d.Mode != modeOwner || g.kind != "desktop" {
		t.Fatalf("%+v kind=%q", d, g.kind)
	}
	if got := g.retryIn(guardInput{Now: at(1), OwnerIdleS: 6}); got != 294 {
		t.Fatalf("300s quiet period, owner idle for 6s: want 294, got %d", got)
	}
	// the owner left: the desktop is quiet but the cooldown has not run out yet
	d = g.step(guardInput{Now: at(100), OwnerIdleS: 100})
	if d.Mode != modeOwner || g.kind != "cooldown" {
		t.Fatalf("%+v kind=%q", d, g.kind)
	}
	if got := g.retryIn(guardInput{Now: at(100), OwnerIdleS: 100}); got != 200 {
		t.Fatalf("want 200, got %d", got)
	}
	d = g.step(guardInput{Now: at(400), OwnerIdleS: 400})
	if d.Mode != modeAvailable || g.kind != "" || g.retryIn(guardInput{Now: at(400), OwnerIdleS: 400}) != 0 {
		t.Fatalf("%+v kind=%q", d, g.kind)
	}
}

func TestRetryIsUnknownWhileAnotherProgramUsesTheGPU(t *testing.T) {
	g := newG()
	in := guardInput{Now: at(0), OwnerIdleS: 3600, ForeignGPU: []string{"blender"}}
	g.step(in)
	g.step(in)
	if g.kind != "gpu" || g.retryIn(in) != 0 {
		t.Fatalf("kind=%q retry=%d", g.kind, g.retryIn(in))
	}
}

func TestPausedHasNoCountdown(t *testing.T) {
	g := newG()
	in := guardInput{Now: at(0), OwnerIdleS: 3600, Paused: true}
	g.step(in)
	g.step(in)
	if g.kind != "paused" || g.retryIn(in) != 0 {
		t.Fatalf("kind=%q retry=%d", g.kind, g.retryIn(in))
	}
}

func TestGPUWorkInOurOwnContainerIsNotTheOwner(t *testing.T) {
	id := strings.Repeat("ab12", 16)
	old := ourContainerIDsFn
	ours.at = time.Time{}
	ourContainerIDsFn = func() []string { return []string{id} }
	t.Cleanup(func() { ourContainerIDsFn = old; ours.at = time.Time{} })
	cg := func(pid int) string {
		switch pid {
		case 1: // a module engine: python inside a container the runner started
			return "0::/system.slice/docker-" + id + ".scope"
		case 2: // somebody else's container
			return "0::/system.slice/docker-" + strings.Repeat("cd34", 16) + ".scope"
		}
		return "0::/user.slice/user-1000.slice/session-3.scope"
	}
	got := otherGPUProcs([]gpuProc{{1, "python"}, {2, "python"}, {3, "blender"}}, cg)
	if len(got) != 2 || got[0] != "python" || got[1] != "blender" {
		t.Fatalf("only the process in our container may be ignored, got %v", got)
	}
	ourContainerIDsFn = func() []string { return nil } // no runner: nothing is recognised
	ours.at = time.Time{}
	if got := otherGPUProcs([]gpuProc{{1, "python"}}, cg); len(got) != 1 {
		t.Fatalf("without the runner the old behaviour must stay, got %v", got)
	}
}