package modelserve

import (
	"context"
	"slices"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)

func TestAContextChangeResizesTheEngineAndTheMemoryAndRestartsTheReplicas(t *testing.T) {
	s, _, models, _, fleet := updSetup(t)
	m := models.models["qwen"]
	m.ContextLength, m.MaxOutput = 4096, maxOutput(4096)
	m.Origin = &catalog.Origin{Repo: "o/r", Files: []catalog.OriginFile{{Path: "q4.gguf", Bytes: 2 << 30}}}
	m.Profiles[0].Args = []string{"--ctx-size", "8192", "--parallel", "2"}
	models.models["qwen"] = m

	got, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 32768}, "admin")
	if err != nil {
		t.Fatal(err)
	}
	if got.Before != 4096 || got.After != 32768 || got.Restarted != 2 || got.Note != "" {
		t.Fatalf("%+v", got)
	}
	after := models.models["qwen"]
	if after.ContextLength != 32768 || after.MaxOutput != maxOutput(32768) {
		t.Errorf("the catalog says %d tokens, %d out", after.ContextLength, after.MaxOutput)
	}
	// Two slots kept: the pool is both.
	if a := after.Profiles[0].Args; !slices.Equal(a, []string{"--ctx-size", "65536", "--parallel", "2"}) {
		t.Errorf("args %v", a)
	}
	// 2 GiB of weights, 64k tokens at 28 KiB each, an engine's share.
	if gb := after.Profiles[0].MinMemoryGB; gb < 4 || gb > 6 {
		t.Errorf("memory %d GB", gb)
	}
	for _, r := range fleet.Replicas() {
		if !strings.Contains(r.ID, "-new") {
			t.Errorf("%s was not started again", r.ID)
		}
	}
}

func TestAContextTheCardsCannotHoldIsRefusedWithTheFigures(t *testing.T) {
	sh := shape{kvPerToken: 128 << 10, weights: 8 << 30}
	p := catalog.Profile{Name: "gpu", Engine: "vllm", GPUClass: "t4-16g", GPUsPerNode: 1}
	if _, _, err := resizeProfile(p, ContextRequest{Tokens: 131072}, sh); err == nil || !strings.Contains(err.Error(), "GB") {
		t.Fatalf("not refused: %v", err)
	}
	next, change, err := resizeProfile(p, ContextRequest{Tokens: 4096}, sh)
	if err != nil {
		t.Fatal(err)
	}
	if !slices.Equal(next.Args, []string{"--max-model-len", "4096"}) || !strings.Contains(change.After, "a card of 1") {
		t.Errorf("%v %+v", next.Args, change)
	}
}

func TestAContextIsRefusedForADeclaredModelOrPastWhatItWasTrainedFor(t *testing.T) {
	s, _, models, _, _ := updSetup(t)
	models.declared["qwen"] = true
	if _, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 8192}, "admin"); err == nil {
		t.Error("a declared model's context was changed")
	}
	models.declared["qwen"] = false
	if _, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 100}, "admin"); err == nil {
		t.Error("100 tokens taken")
	}
	if _, err := s.SetContext(context.Background(), "nope", ContextRequest{Tokens: 8192}, "admin"); err == nil {
		t.Error("no such model, and no error")
	}
}

func TestEngineArgumentsAreSetInEitherSpelling(t *testing.T) {
	if a := setArg([]string{"--ctx-size=4096", "-ngl", "99"}, "--ctx-size", "8192"); !slices.Equal(a, []string{"--ctx-size=8192", "-ngl", "99"}) {
		t.Error(a)
	}
	if a := setArg([]string{"-ngl", "99"}, "--parallel", "4"); !slices.Equal(a, []string{"-ngl", "99", "--parallel", "4"}) {
		t.Error(a)
	}
	if argInt([]string{"--parallel=3"}, "--parallel") != 3 || argInt(nil, "--parallel") != 0 {
		t.Error("argInt")
	}
	p := catalog.Profile{Name: "x", Engine: "comfyui"}
	if _, _, err := resizeProfile(p, ContextRequest{Tokens: 4096}, shape{}); err == nil {
		t.Error("a comfyui profile took a context")
	}
}

func TestTheLongestAnswerAloneChangesOnlyTheCatalogAndStartsNothing(t *testing.T) {
	s, _, models, _, fleet := updSetup(t)
	m := models.models["qwen"]
	m.ContextLength, m.MaxOutput = 4096, 1024
	m.Profiles[0].Args = []string{"--ctx-size", "4096"}
	models.models["qwen"] = m
	before := fleet.Replicas()

	got, err := s.SetContext(context.Background(), "qwen", ContextRequest{MaxOutput: 3000}, "admin")
	if err != nil {
		t.Fatal(err)
	}
	if !got.OutputOnly || got.MaxOutput != 3000 || got.MaxOutputBefore != 1024 || got.Restarted != 0 {
		t.Fatalf("%+v", got)
	}
	after := models.models["qwen"]
	if after.ContextLength != 4096 || after.MaxOutput != 3000 || !slices.Equal(after.Profiles[0].Args, []string{"--ctx-size", "4096"}) {
		t.Errorf("catalog %d/%d, args %v", after.ContextLength, after.MaxOutput, after.Profiles[0].Args)
	}
	for i, r := range fleet.Replicas() {
		if r.ID != before[i].ID {
			t.Errorf("%s was started again for a limit no engine reads", before[i].ID)
		}
	}
	// An answer longer than the context is refused, alone or with one.
	if _, err := s.SetContext(context.Background(), "qwen", ContextRequest{MaxOutput: 5000}, "admin"); err == nil {
		t.Error("an answer longer than the context was accepted")
	}
	if _, err := s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 8192, MaxOutput: 9000}, "admin"); err == nil {
		t.Error("an answer longer than the new context was accepted")
	}
	// With a context, the one asked for wins over the one that follows it.
	got, err = s.SetContext(context.Background(), "qwen", ContextRequest{Tokens: 8192, MaxOutput: 6000}, "admin")
	if err != nil || got.MaxOutput != 6000 || models.models["qwen"].MaxOutput != 6000 {
		t.Errorf("%+v %v", got, err)
	}
}

func TestStrataTakesItsContextInItsSetupArguments(t *testing.T) {
	p := catalog.Profile{Name: "one-card", Engine: "strata", GPUClass: "6000-24g",
		Args: []string{"--family", "qwen", "--size", "IQ2_XS"}}
	got, change, err := resizeProfile(p, ContextRequest{Tokens: 32768}, shape{kvPerToken: 1 << 20, weights: 40 << 30})
	if err != nil {
		t.Fatal(err)
	}
	// Not refused for weights larger than the card: its experts live in RAM.
	if !slices.Equal(got.Args, []string{"--family", "qwen", "--size", "IQ2_XS", "--context", "32768"}) || !slices.Equal(change.Args, got.Args) {
		t.Errorf("args %v", got.Args)
	}
	if n := contextOf(catalog.Model{ContextLength: 4096}, got); n != 32768 {
		t.Errorf("contextOf %d", n)
	}
}