package agentsvc

import (
	"slices"
	"testing"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

func TestANodeIsAskedOnlyForEnginesItCanInstall(t *testing.T) {
	// An agent older than the list says nothing, and could install llama.cpp
	// alone. Asked for vLLM, it showed as failing until it was upgraded.
	if got := installable(&llmgwv1.Register{}); !slices.Equal(got, []string{"llamacpp"}) {
		t.Errorf("an agent that said nothing: %v", got)
	}
	// This agent says what it can, and that is vLLM and SGLang besides.
	now := &llmgwv1.Register{InstallableEngines: provision.Installable()}
	for _, e := range []string{"llamacpp", "vllm", "sglang"} {
		if !slices.Contains(installable(now), e) {
			t.Errorf("this agent does not offer %s: %v", e, installable(now))
		}
	}
	if slices.Contains(installable(now), "uv") {
		t.Error("uv, a tool for the others, is offered as an engine")
	}
}

func TestEachNodeIsToldTheEnginesDecidedForIt(t *testing.T) {
	s := &Service{}
	s.SetEnginesFor(func(n inventory.Node) map[string]string {
		if n.ID == "volta" {
			return map[string]string{"llamacpp": "b1", "vllm": "0.25.1"}
		}
		return map[string]string{"llamacpp": "b1", "vllm": "0.30.0"}
	})
	reg := &llmgwv1.Register{InstallableEngines: []string{"llamacpp", "vllm"}}
	want := func(p *llmgwv1.Provision) map[string]string {
		out := map[string]string{}
		for _, e := range p.GetEngines() {
			out[e.GetEngine()] = e.GetVersion()
		}
		return out
	}
	if got := want(s.provisionFor(inventory.Node{ID: "volta"}, installable(reg))); got["vllm"] != "0.25.1" {
		t.Errorf("the Volta node: %v", got)
	}
	if got := want(s.provisionFor(inventory.Node{ID: "hopper"}, installable(reg))); got["vllm"] != "0.30.0" {
		t.Errorf("the Hopper node: %v", got)
	}
	// A node that disabled vLLM does not offer it, and is not asked for it.
	if got := want(s.provisionFor(inventory.Node{ID: "hopper"}, []string{"llamacpp"})); got["vllm"] != "" {
		t.Errorf("asked a node for an engine it does not install: %v", got)
	}
}

func TestAnOfflineClusterTellsNodesWhereOnTheMirrorAndWhatEachFileMustBe(t *testing.T) {
	s := &Service{}
	s.SetEngines(map[string]string{"llamacpp": "b1"})
	if w := s.provisionFor(inventory.Node{ID: "n"}, []string{"llamacpp"}).GetEngines()[0]; w.GetOfflineBase() != "" {
		t.Fatalf("an online cluster sent a mirror: %v", w)
	}
	s.SetOffline(func(engine, version string) (string, map[string]string) {
		return "http://mirror:8081/_llmgw/engines/" + engine + "/" + version,
			map[string]string{"llama-" + version + "-bin-ubuntu-x64.tar.gz": "abc"}
	})
	w := s.provisionFor(inventory.Node{ID: "n"}, []string{"llamacpp"}).GetEngines()[0]
	if w.GetOfflineBase() != "http://mirror:8081/_llmgw/engines/llamacpp/b1" ||
		w.GetOfflineSha256()["llama-b1-bin-ubuntu-x64.tar.gz"] != "abc" {
		t.Errorf("the node was told %v", w)
	}
}

func TestABuildChosenForANodeIsWhatItIsTold(t *testing.T) {
	s := &Service{}
	s.SetEngines(map[string]string{"llamacpp": "b1", "vllm": "0.30.0"})
	s.SetEngineChoices(func(node string) map[string]string {
		return map[string]map[string]string{"gpu07": {"vllm": "cu129"}, "gpu08": {"vllm": "none"}}[node]
	})
	can := []string{"llamacpp", "vllm"}
	variant := func(node string) (string, bool) {
		for _, e := range s.provisionFor(inventory.Node{ID: node}, can).GetEngines() {
			if e.GetEngine() == "vllm" {
				return e.GetVariant(), true
			}
		}
		return "", false
	}
	if v, ok := variant("gpu07"); !ok || v != "cu129" {
		t.Errorf("gpu07: %q, %v", v, ok)
	}
	if _, ok := variant("gpu08"); ok {
		t.Error("gpu08 was asked for a vLLM it was chosen to run none of")
	}
	if v, ok := variant("gpu09"); !ok || v != "" {
		t.Errorf("gpu09, left to its hardware: %q, %v", v, ok)
	}
}