package catalog

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

func model(name string, profiles ...Profile) Model {
	if len(profiles) == 0 {
		profiles = []Profile{{Name: "p1", GPUClass: "h100-80g", TP: 4}}
	}
	return Model{Name: name, Source: "/models/" + name, Profiles: profiles}
}

func TestResolveNameAndAlias(t *testing.T) {
	m := model("llama-3.3-70b")
	m.Aliases = []string{"llama", "llama3"}
	c, err := New([]Model{m})
	if err != nil {
		t.Fatal(err)
	}

	for _, name := range []string{"llama-3.3-70b", "llama", "llama3"} {
		got, ok := c.Resolve(name)
		if !ok {
			t.Fatalf("%q did not resolve", name)
		}
		if got.Name != "llama-3.3-70b" {
			t.Errorf("%q resolved to %q", name, got.Name)
		}
	}
	if _, ok := c.Resolve("nope"); ok {
		t.Error("an unknown name resolved")
	}
}

func TestAliasCollisionRejected(t *testing.T) {
	// An alias shadowing another model would misroute every request for it,
	// silently, at request time.
	a := model("qwen")
	b := model("llama")
	b.Aliases = []string{"qwen"}

	_, err := New([]Model{a, b})
	if err == nil || !strings.Contains(err.Error(), "collides") {
		t.Errorf("collision not rejected: %v", err)
	}
}

func TestDuplicateAliasRejected(t *testing.T) {
	a := model("a")
	a.Aliases = []string{"shared"}
	b := model("b")
	b.Aliases = []string{"shared"}

	_, err := New([]Model{a, b})
	if err == nil || !strings.Contains(err.Error(), "claimed by both") {
		t.Errorf("duplicate alias not rejected: %v", err)
	}
}

func TestUnknownFallbackRejected(t *testing.T) {
	// A typo in a fallback should stop a deployment, not surface during the
	// outage when the fallback is finally needed.
	m := model("llama")
	m.Fallbacks = []string{"qwen-typo"}

	_, err := New([]Model{m})
	if err == nil || !strings.Contains(err.Error(), "fallback") {
		t.Errorf("unknown fallback not rejected: %v", err)
	}
}

func TestFallbackToAliasAccepted(t *testing.T) {
	a := model("qwen2.5-32b")
	a.Aliases = []string{"qwen"}
	b := model("llama")
	b.Fallbacks = []string{"qwen"}

	if _, err := New([]Model{a, b}); err != nil {
		t.Errorf("a fallback naming an alias was rejected: %v", err)
	}
}

func TestParallelismArithmeticValidated(t *testing.T) {
	// Wrong arithmetic hands the engine a spec it cannot satisfy, and it fails
	// minutes later during model loading rather than at config load.
	_, err := New([]Model{model("m", Profile{
		Name: "bad", GPUClass: "h100-80g", TP: 4, Nodes: 2, GPUsPerNode: 8,
	})})
	if err == nil || !strings.Contains(err.Error(), "does not match") {
		t.Errorf("bad arithmetic accepted: %v", err)
	}

	// 2 nodes x 8 GPUs = TP 16 is consistent.
	if _, err := New([]Model{model("m", Profile{
		Name: "ok", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
	})}); err != nil {
		t.Errorf("valid arithmetic rejected: %v", err)
	}
}

func TestDefaults(t *testing.T) {
	c, err := New([]Model{model("m", Profile{Name: "p", GPUClass: "h100-80g", TP: 4})})
	if err != nil {
		t.Fatal(err)
	}
	m, _ := c.Resolve("m")
	p := m.Profiles[0]

	if p.Engine != "vllm" {
		t.Errorf("engine=%q", p.Engine)
	}
	if p.Nodes != 1 {
		t.Errorf("nodes=%d", p.Nodes)
	}
	if p.GPUsPerNode != 4 {
		t.Errorf("gpus_per_node=%d; a single-node profile should default to TP", p.GPUsPerNode)
	}
	if p.StartupTimeout != 20*time.Minute {
		t.Errorf("startup_timeout=%v", p.StartupTimeout)
	}
	if p.Weight <= 0 {
		t.Errorf("weight=%v", p.Weight)
	}
	// Max is no longer filled in for a model that names no minimum. With min 0
	// it now means "configured and deliberately not running", and defaulting it
	// to 1 would make every model without a scaling stanza loadable on demand —
	// so any client request could start a 70B nobody intended to run.
	if m.Scaling.Min == 0 && m.Scaling.Max != 0 {
		t.Errorf("scaling max=%d; an unscaled model should stay off, not become on-demand", m.Scaling.Max)
	}
	if m.Scaling.OnDemand() {
		t.Error("a model with no scaling stanza was made on-demand")
	}
}

func TestMaxDefaultsToTheMinimumWhenOneIsGiven(t *testing.T) {
	// A model that says it wants two replicas and nothing about a maximum wants
	// two, not one and not unbounded.
	built, err := New([]Model{{
		Name: "m", Source: "/m", Scaling: Scaling{Min: 2},
		Profiles: []Profile{{Name: "p", Engine: "llamacpp"}},
	}})
	if err != nil {
		t.Fatal(err)
	}
	m, _ := built.Resolve("m")
	if m.Scaling.Max != 2 {
		t.Errorf("max=%d, want 2", m.Scaling.Max)
	}
	if m.Scaling.OnDemand() {
		t.Error("a model with min 2 was reported on-demand")
	}
}

func TestOnDemandNeedsBothZeroMinimumAndARealMaximum(t *testing.T) {
	for _, tc := range []struct {
		name string
		s    Scaling
		want bool
	}{
		{"off", Scaling{Min: 0, Max: 0}, false},
		{"on demand", Scaling{Min: 0, Max: 1}, true},
		{"always on", Scaling{Min: 1, Max: 1}, false},
		{"always on, room to grow", Scaling{Min: 1, Max: 4}, false},
	} {
		if got := tc.s.OnDemand(); got != tc.want {
			t.Errorf("%s: OnDemand() = %v, want %v", tc.name, got, tc.want)
		}
	}
}

func TestRequiredFields(t *testing.T) {
	cases := map[string]Model{
		"no name":     {Source: "/m", Profiles: []Profile{{Name: "p", GPUClass: "c", TP: 1}}},
		"no source":   {Name: "m", Profiles: []Profile{{Name: "p", GPUClass: "c", TP: 1}}},
		"no profiles": {Name: "m", Source: "/m"},
		"no class":    {Name: "m", Source: "/m", Profiles: []Profile{{Name: "p", TP: 1}}},
	}
	for name, m := range cases {
		t.Run(name, func(t *testing.T) {
			if _, err := New([]Model{m}); err == nil {
				t.Error("accepted")
			}
		})
	}
}

func TestScalingBounds(t *testing.T) {
	m := model("m")
	m.Scaling = Scaling{Min: 5, Max: 2}
	if _, err := New([]Model{m}); err == nil {
		t.Error("max below min was accepted")
	}
}

func TestDuplicateProfileRejected(t *testing.T) {
	_, err := New([]Model{model("m",
		Profile{Name: "p", GPUClass: "h100-80g", TP: 1},
		Profile{Name: "p", GPUClass: "a100-40g", TP: 1},
	)})
	if err == nil || !strings.Contains(err.Error(), "twice") {
		t.Errorf("duplicate profile accepted: %v", err)
	}
}

func TestPlacementProfilesCarryPriority(t *testing.T) {
	c, _ := New([]Model{model("m",
		Profile{Name: "fast", GPUClass: "h100-80g", TP: 4, Priority: 100},
		Profile{Name: "slow", GPUClass: "a100-40g", TP: 4, Priority: 50},
	)})
	m, _ := c.Resolve("m")

	pp := m.PlacementProfiles()
	if len(pp) != 2 {
		t.Fatalf("got %d placement profiles", len(pp))
	}
	byName := map[string]int{}
	for _, p := range pp {
		byName[p.Name] = p.Priority
	}
	if byName["fast"] != 100 || byName["slow"] != 50 {
		t.Errorf("priorities lost in conversion: %v", byName)
	}
}

func TestEngineProfileCarriesModelIdentity(t *testing.T) {
	// The served name is what clients ask for and need not match the path.
	c, _ := New([]Model{model("llama-3.3-70b")})
	m, _ := c.Resolve("llama-3.3-70b")

	ep, ok := m.EngineProfile("p1")
	if !ok {
		t.Fatal("engine profile not found")
	}
	if ep.ModelPath != "/models/llama-3.3-70b" {
		t.Errorf("ModelPath=%q", ep.ModelPath)
	}
	if ep.ServedModelName != "llama-3.3-70b" {
		t.Errorf("ServedModelName=%q", ep.ServedModelName)
	}

	if _, ok := m.EngineProfile("nope"); ok {
		t.Error("an unknown profile resolved")
	}
}

func TestEngineProfileArgsAreCopied(t *testing.T) {
	// The reconciler hands these to an engine adapter; mutating them must not
	// corrupt the catalog for every later replica.
	c, _ := New([]Model{model("m", Profile{
		Name: "p", GPUClass: "h100-80g", TP: 1, Args: []string{"--max-model-len", "32768"},
	})})
	m, _ := c.Resolve("m")

	ep, _ := m.EngineProfile("p")
	ep.Args[0] = "tampered"

	again, _ := m.EngineProfile("p")
	if again.Args[0] != "--max-model-len" {
		t.Error("EngineProfile exposed the catalog's own slice")
	}
}

func TestNamesAndModelsOrder(t *testing.T) {
	c, _ := New([]Model{model("zeta"), model("alpha")})

	// Models keeps declaration order, so configuration reads as written.
	if got := c.Models(); got[0].Name != "zeta" {
		t.Errorf("Models order=%q", got[0].Name)
	}
	// Names is sorted, for stable CLI output.
	if got := c.Names(); got[0] != "alpha" {
		t.Errorf("Names order=%v", got)
	}
}

func TestSimulatedNeedsEveryProfileToBeSimulated(t *testing.T) {
	// The question the flag answers is "can this request come back as real
	// inference?". One real profile means yes, however many sim profiles sit
	// beside it, so a mixed model must not be labelled simulated.
	for _, tc := range []struct {
		name    string
		engines []string
		wantSim bool
	}{
		{"all simulated", []string{"sim", "sim"}, true},
		{"one real profile", []string{"sim", "llamacpp"}, false},
		{"no simulation", []string{"llamacpp"}, false},
		{"no profiles at all", nil, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var m Model
			for _, e := range tc.engines {
				m.Profiles = append(m.Profiles, Profile{Engine: e})
			}
			if got := m.Simulated(); got != tc.wantSim {
				t.Errorf("Simulated() = %v, want %v", got, tc.wantSim)
			}
		})
	}
}

func TestAProfileCannotSetWhatRunsCodeInTheEngine(t *testing.T) {
	// A profile's env reaches an engine started as root. LD_PRELOAD from a
	// model definition would make "add a model" into "run this everywhere".
	for _, k := range []string{"LD_PRELOAD", "ld_library_path", "PYTHONSTARTUP", "PATH", "BASH_ENV"} {
		m := Model{Name: "m", Source: "http://x/m.gguf", Profiles: []Profile{{
			Name: "cpu", Engine: "llamacpp", Env: map[string]string{k: "/tmp/x"},
		}}}
		if _, err := New([]Model{m}); err == nil || !strings.Contains(err.Error(), "not allowed") {
			t.Errorf("%s: %v", k, err)
		}
	}
	ok := Model{Name: "m", Source: "http://x/m.gguf", Profiles: []Profile{{
		Name: "cpu", Engine: "llamacpp", Env: map[string]string{"OMP_NUM_THREADS": "4", "PYTHONUNBUFFERED": "1"},
	}}}
	if _, err := New([]Model{ok}); err != nil {
		t.Errorf("ordinary tuning refused: %v", err)
	}
}

func TestALlamaCppReplicaOnSeveralGPUsSaysHowManyAsGPUsPerNode(t *testing.T) {
	four := model("big", Profile{Name: "t4", Engine: "llamacpp", GPUClass: "t4", GPUsPerNode: 4})
	c, err := New([]Model{four})
	if err != nil {
		t.Fatalf("gpus_per_node 4 on llama.cpp: %v", err)
	}
	m, _ := c.Resolve("big")
	if pl := m.PlacementProfiles(); len(pl) != 1 || pl[0].GPUsPerNode != 4 {
		t.Errorf("placement was not asked for four GPUs: %+v", pl)
	}
	// tp, written before gpus_per_node meant anything for llama.cpp, is read
	// as the same width rather than refused: refusing dropped the model on
	// the upgrade that made the width work.
	old := model("old", Profile{Name: "t4", Engine: "llamacpp", GPUClass: "t4", TP: 4})
	c, err = New([]Model{old})
	if err != nil {
		t.Fatalf("tp 4 on llama.cpp: %v", err)
	}
	got, _ := c.Resolve("old")
	if p := got.PlacementProfiles()[0]; p.GPUsPerNode != 4 {
		t.Errorf("tp 4 should place on four GPUs: %+v", p)
	}
	// Both given and disagreeing is ambiguous, and refused; so is more than
	// one node.
	for _, p := range []Profile{
		{Name: "both", Engine: "llamacpp", GPUClass: "t4", TP: 4, GPUsPerNode: 2},
		{Name: "nodes", Engine: "llamacpp", GPUClass: "t4", Nodes: 2, GPUsPerNode: 1},
	} {
		if _, err := New([]Model{model("x", p)}); err == nil {
			t.Errorf("%s: accepted", p.Name)
		}
	}
}

func TestARepositoryPageIsReadAsTheRepositoryForVLLM(t *testing.T) {
	// The field: an AWQ model's page given as a vLLM source, fetched as one
	// file of HTML. vLLM downloads a repository id itself.
	for _, src := range []string{
		"https://huggingface.co/Qwen/Qwen2.5-32B-Instruct-AWQ",
		"https://huggingface.co/Qwen/Qwen2.5-32B-Instruct-AWQ/tree/main",
	} {
		m := model("awq")
		m.Source = src
		c, err := New([]Model{m})
		if err != nil {
			t.Fatalf("%s: %v", src, err)
		}
		if got, _ := c.Resolve("awq"); got.Source != "Qwen/Qwen2.5-32B-Instruct-AWQ" {
			t.Errorf("%s became %q", src, got.Source)
		}
	}
	// llama.cpp needs a file, which a page is not.
	g := model("gguf", Profile{Name: "cpu", Engine: "llamacpp", MinMemoryGB: 4})
	g.Source = "https://huggingface.co/Qwen/X-GGUF"
	if _, err := New([]Model{g}); err == nil || !strings.Contains(err.Error(), "needs a .gguf") {
		t.Errorf("a page for llama.cpp: %v", err)
	}
	f := model("file")
	f.Source = "https://huggingface.co/Qwen/X-GGUF/resolve/abc123/x.gguf"
	if _, err := New([]Model{f}); err != nil {
		t.Errorf("a file under /resolve/ is a source: %v", err)
	}
}

func TestAModelThatWouldRunCodeOnItsNodesIsRefusedWhereItIsDefined(t *testing.T) {
	for _, p := range []Profile{
		{Name: "p", Engine: "vllm", GPUClass: "h100-80g", TP: 1, Args: []string{"--trust-remote-code"}},
		{Name: "p", Engine: "llamacpp", MinMemoryGB: 4, Args: []string{"--log-file", "/etc/cron.d/x"}},
		{Name: "p", Engine: "vllm", GPUClass: "h100-80g", TP: 1, Env: map[string]string{"NCCL_NET_PLUGIN": "/tmp/x.so"}},
		{Name: "p", Engine: "vllm", GPUClass: "h100-80g", TP: 1, Env: map[string]string{"CUDA_INJECTION64_PATH": "/tmp/x.so"}},
	} {
		if _, err := New([]Model{model("m", p)}); err == nil {
			t.Errorf("accepted %v %v", p.Args, p.Env)
		}
	}
	ok := Profile{Name: "p", Engine: "llamacpp", MinMemoryGB: 4, Args: []string{"--ctx-size", "16384"},
		Env: map[string]string{"OMP_NUM_THREADS": "8"}}
	if _, err := New([]Model{model("m", ok)}); err != nil {
		t.Errorf("a tuned model was refused: %v", err)
	}
}

func TestAGuardModelSaysHowItIsReadAndAlwaysRuns(t *testing.T) {
	base := func() Model {
		return Model{Name: "lg", Source: "/m/lg.gguf", Role: RoleGuard, Guard: &GuardModel{Format: "llamaguard3"},
			Profiles: []Profile{{Name: "cpu", Engine: "llamacpp"}}, Scaling: Scaling{Min: 1, Max: 2}}
	}
	if err := base().validate(); err != nil {
		t.Fatal(err)
	}
	for name, mut := range map[string]func(*Model){
		"no format":      func(m *Model) { m.Guard = nil },
		"bad format":     func(m *Model) { m.Guard.Format = "vibes" },
		"on demand":      func(m *Model) { m.Scaling = Scaling{Min: 0, Max: 1} },
		"bad role":       func(m *Model) { m.Role = "judge" },
		"guard w/o role": func(m *Model) { m.Role = "" },
	} {
		m := base()
		mut(&m)
		if m.validate() == nil {
			t.Errorf("%s accepted", name)
		}
	}
	if base().EffectiveClass() != ClassCritical {
		t.Error("a judge is not critical by default")
	}
}