package main

import (
	"context"
	"encoding/json"
	"testing"
	"time"
)

type fakeBackend struct {
	loadedM    map[string]int64
	installedM map[string]int64
	total      uint64
	avail      uint64
	unloaded   []string
}

func (f *fakeBackend) loaded() (map[string]int64, error) {
	out := map[string]int64{}
	for k, v := range f.loadedM {
		out[k] = v
	}
	return out, nil
}
func (f *fakeBackend) installed() (map[string]int64, error) { return f.installedM, nil }
func (f *fakeBackend) unload(n string) error {
	f.unloaded = append(f.unloaded, n)
	if sz, ok := f.loadedM[n]; ok {
		f.avail += uint64(sz)
		delete(f.loadedM, n)
	}
	return nil
}
func (f *fakeBackend) memory() (uint64, uint64) { return f.total, f.avail }

const gib = int64(1 << 30)

func newTestPool(fb *fakeBackend) (*pool, *time.Time) {
	p := newPool(fb)
	clock := time.Unix(1000, 0)
	p.now = func() time.Time { return clock }
	p.sleep = func(d time.Duration) { clock = clock.Add(d) } // advance fake time instead of sleeping
	p.reserveBytes = 8 << 30
	p.headroom = 0
	return p, &clock
}

func TestLoadedModelIsReusedWithoutEviction(t *testing.T) {
	fb := &fakeBackend{loadedM: map[string]int64{"a:latest": 10 * gib}, installedM: map[string]int64{"a:latest": 10 * gib}, total: 100 << 30, avail: 50 << 30}
	p, _ := newTestPool(fb)
	rel, err := p.acquire(context.Background(), "a")
	if err != nil {
		t.Fatal(err)
	}
	rel()
	if len(fb.unloaded) != 0 {
		t.Fatalf("unexpected eviction: %v", fb.unloaded)
	}
}

func TestUnknownModelIs404(t *testing.T) {
	fb := &fakeBackend{installedM: map[string]int64{}, total: 100 << 30, avail: 90 << 30}
	p, _ := newTestPool(fb)
	_, err := p.acquire(context.Background(), "nope")
	if err == nil || err.Status != 404 {
		t.Fatalf("want 404, got %v", err)
	}
}

func TestFitsWithoutEviction(t *testing.T) {
	fb := &fakeBackend{loadedM: map[string]int64{"a:latest": 10 * gib}, installedM: map[string]int64{"a:latest": 10 * gib, "b:latest": 10 * gib}, total: 100 << 30, avail: 60 << 30}
	p, _ := newTestPool(fb)
	rel, err := p.acquire(context.Background(), "b")
	if err != nil {
		t.Fatal(err)
	}
	rel()
	if len(fb.unloaded) != 0 {
		t.Fatalf("evicted although memory was enough: %v", fb.unloaded)
	}
}

func TestEvictsIdleLeastRecentlyUsedFirst(t *testing.T) {
	fb := &fakeBackend{
		loadedM:    map[string]int64{"old:latest": 20 * gib, "recent:latest": 20 * gib},
		installedM: map[string]int64{"old:latest": 20 * gib, "recent:latest": 20 * gib, "new:latest": 30 * gib},
		total:      64 << 30, avail: 12 << 30, // 12 free - 8 reserve = 4 usable; needs 30*1.125
	}
	p, clock := newTestPool(fb)
	p.lastUse["old:latest"] = clock.Add(-10 * time.Minute)
	p.lastUse["recent:latest"] = clock.Add(-1 * time.Minute)
	p.seenAt["old:latest"] = clock.Add(-time.Hour)
	p.seenAt["recent:latest"] = clock.Add(-time.Hour)
	rel, err := p.acquire(context.Background(), "new")
	if err != nil {
		t.Fatal(err)
	}
	rel()
	if len(fb.unloaded) == 0 || fb.unloaded[0] != "old:latest" {
		t.Fatalf("expected old:latest evicted first, got %v", fb.unloaded)
	}
}

func TestNeverEvictsBusyModelAndEventually503(t *testing.T) {
	fb := &fakeBackend{
		loadedM:    map[string]int64{"busy:latest": 40 * gib},
		installedM: map[string]int64{"busy:latest": 40 * gib, "new:latest": 30 * gib},
		total:      64 << 30, avail: 12 << 30,
	}
	p, clock := newTestPool(fb)
	p.seenAt["busy:latest"] = clock.Add(-time.Hour)
	p.inflight["busy:latest"] = 1
	p.maxWait = 3 * time.Second
	_, err := p.acquire(context.Background(), "new")
	if err == nil || err.Status != 503 {
		t.Fatalf("want 503, got %v", err)
	}
	if len(fb.unloaded) != 0 {
		t.Fatalf("busy model was evicted: %v", fb.unloaded)
	}
}

func TestWaitsThenSucceedsWhenBusyModelFinishes(t *testing.T) {
	fb := &fakeBackend{
		loadedM:    map[string]int64{"busy:latest": 40 * gib},
		installedM: map[string]int64{"busy:latest": 40 * gib, "new:latest": 30 * gib},
		total:      64 << 30, avail: 12 << 30,
	}
	p, clock := newTestPool(fb)
	p.seenAt["busy:latest"] = clock.Add(-time.Hour)
	p.inflight["busy:latest"] = 1
	n := 0
	p.sleep = func(d time.Duration) { // after two polls the busy request finishes
		*clock = clock.Add(d)
		n++
		if n == 2 {
			p.mu.Lock()
			p.inflight["busy:latest"] = 0
			p.mu.Unlock()
		}
	}
	rel, err := p.acquire(context.Background(), "new")
	if err != nil {
		t.Fatalf("should have succeeded after the busy model finished: %v", err)
	}
	rel()
	if len(fb.unloaded) != 1 || fb.unloaded[0] != "busy:latest" {
		t.Fatalf("expected the now-idle model evicted, got %v", fb.unloaded)
	}
}

func TestFreshlyLoadedModelIsProtectedByMinResidency(t *testing.T) {
	fb := &fakeBackend{
		loadedM:    map[string]int64{"fresh:latest": 40 * gib},
		installedM: map[string]int64{"fresh:latest": 40 * gib, "new:latest": 30 * gib},
		total:      64 << 30, avail: 12 << 30,
	}
	p, clock := newTestPool(fb)
	p.seenAt["fresh:latest"] = *clock // just loaded
	p.maxWait = time.Second
	_, err := p.acquire(context.Background(), "new")
	if err == nil || err.Status != 503 {
		t.Fatalf("fresh model should not be evicted immediately: %v", err)
	}
	if len(fb.unloaded) != 0 {
		t.Fatalf("fresh model evicted: %v", fb.unloaded)
	}
}

func TestTwoConcurrentLoadsCannotOverCommit(t *testing.T) {
	fb := &fakeBackend{
		installedM: map[string]int64{"x:latest": 30 * gib, "y:latest": 30 * gib},
		total:      64 << 30, avail: 50 << 30, // room for one 30 GB model after the 8 GB reserve, not two
	}
	p, _ := newTestPool(fb)
	p.maxWait = time.Second
	relX, err := p.acquire(context.Background(), "x")
	if err != nil {
		t.Fatal(err)
	}
	defer relX()
	_, err = p.acquire(context.Background(), "y")
	if err == nil || err.Status != 503 {
		t.Fatalf("second model must not be admitted while the first is still loading: %v", err)
	}
}

func TestTooLargeModelIsRejectedImmediately(t *testing.T) {
	fb := &fakeBackend{installedM: map[string]int64{"huge:latest": 200 * gib}, total: 128 << 30, avail: 120 << 30}
	p, _ := newTestPool(fb)
	_, err := p.acquire(context.Background(), "huge")
	if err == nil || err.Status != 503 {
		t.Fatalf("want 503 for a model larger than the machine: %v", err)
	}
}

func TestUnknownMemoryDoesNotBlock(t *testing.T) {
	fb := &fakeBackend{installedM: map[string]int64{"a:latest": 5 * gib}} // total == 0
	p, _ := newTestPool(fb)
	rel, err := p.acquire(context.Background(), "a")
	if err != nil {
		t.Fatal(err)
	}
	rel()
}

func TestResolveModelUsesTheInstalledName(t *testing.T) {
	inst := map[string]int64{"qwen2.5:0.5b": 1, "qwen3-0.6b": 1, "nvidia/Qwen3.8-27B-NVFP4": 1, "imp:latest": 1}
	cases := map[string]string{
		"qwen2.5:0.5b":             "qwen2.5:0.5b",
		"qwen3-0.6b":               "qwen3-0.6b", // LM Studio key: must NOT become qwen3-0.6b:latest
		"nvidia/Qwen3.8-27B-NVFP4": "nvidia/Qwen3.8-27B-NVFP4",
		"imp":                      "imp:latest", // Ollama convention
		"unknown":                  "unknown",
	}
	for in, want := range cases {
		if got := resolveModel(in, inst); got != want {
			t.Errorf("resolveModel(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestAcquireAcceptsPlainNamesFromAnyRuntime(t *testing.T) {
	fb := &fakeBackend{installedM: map[string]int64{"qwen3-0.6b": 1 * gib}, total: 100 << 30, avail: 90 << 30}
	p, _ := newTestPool(fb)
	rel, err := p.acquire(context.Background(), "qwen3-0.6b")
	if err != nil {
		t.Fatalf("a plain LM Studio style name must be accepted: %v", err)
	}
	rel()
}

func TestWithModelRewritesOnlyTheModelField(t *testing.T) {
	out := withModel([]byte(`{"model":"qwen","messages":[{"role":"user","content":"hi"}],"stream":true}`), "qwen:latest")
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil || m["model"] != "qwen:latest" || m["stream"] != true || m["messages"] == nil {
		t.Fatalf("%s %v", out, err)
	}
}

func TestFreeUpUnloadsOnlyAsManyIdleModelsAsNeededLeastRecentlyUsedFirst(t *testing.T) {
	fb := &fakeBackend{
		loadedM: map[string]int64{"old:latest": 20 * gib, "mid:latest": 20 * gib, "new:latest": 20 * gib, "busy:latest": 20 * gib},
		total:   128 << 30, avail: 20 << 30,
	}
	p, clock := newTestPool(fb)
	p.lastUse["old:latest"] = clock.Add(-3 * time.Hour)
	p.lastUse["mid:latest"] = clock.Add(-2 * time.Hour)
	p.lastUse["new:latest"] = clock.Add(-time.Minute)
	p.lastUse["busy:latest"] = clock.Add(-5 * time.Hour)
	p.inflight["busy:latest"] = 1 // in use right now: never touched, however long ago it was last used

	// needs 38 GB; 20 free, 8 kept for the owner: 12 usable, so 26 more are needed = two models (the two least recently used)
	ok, gone := p.freeUp(uint64(38) << 30)
	if !ok {
		t.Fatal("there is room to be made")
	}
	if len(gone) != 2 || gone[0] != "old:latest" || gone[1] != "mid:latest" {
		t.Fatalf("only the two least recently used idle models may go, in that order: %v", gone)
	}
	if _, still := fb.loadedM["new:latest"]; !still {
		t.Error("a model used a minute ago was unloaded")
	}
	if _, still := fb.loadedM["busy:latest"]; !still {
		t.Error("a model with a request in flight was unloaded")
	}
}

func TestFreeUpDoesNothingWhenThereIsRoom(t *testing.T) {
	fb := &fakeBackend{loadedM: map[string]int64{"a:latest": 20 * gib}, total: 128 << 30, avail: 80 << 30}
	p, _ := newTestPool(fb)
	ok, gone := p.freeUp(uint64(24) << 30)
	if !ok || len(gone) != 0 || len(fb.unloaded) != 0 {
		t.Fatalf("models stay while there is room: ok=%v gone=%v", ok, gone)
	}
}

func TestFreeUpSaysSoWhenItCannotMakeEnoughRoom(t *testing.T) {
	fb := &fakeBackend{loadedM: map[string]int64{"busy:latest": 60 * gib, "idle:latest": 5 * gib}, total: 128 << 30, avail: 10 << 30}
	p, _ := newTestPool(fb)
	p.inflight["busy:latest"] = 1
	ok, gone := p.freeUp(uint64(60) << 30)
	if ok {
		t.Error("it cannot claim room it did not make")
	}
	if len(gone) != 1 || gone[0] != "idle:latest" {
		t.Errorf("it should still have tried what it could: %v", gone)
	}
}