package modelserve

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
)

// ggufHeader is enough of a GGUF file for its architecture and shape.
func ggufHeader() []byte {
	var b bytes.Buffer
	w := func(v any) { _ = binary.Write(&b, binary.LittleEndian, v) }
	str := func(s string) { w(uint64(len(s))); b.WriteString(s) }
	u32 := func(k string, v uint32) { str(k); w(uint32(4)); w(v) }
	b.WriteString("GGUF")
	w(uint32(3))
	w(uint64(0))
	w(uint64(5))
	str("general.architecture")
	w(uint32(8))
	str("qwen2")
	u32("qwen2.block_count", 28)
	u32("qwen2.attention.head_count", 12)
	u32("qwen2.attention.head_count_kv", 2)
	u32("qwen2.embedding_length", 1536)
	return b.Bytes()
}

type fakeRegistry struct {
	repo hf.Repo
	base string
}

func (f fakeRegistry) Base() string { return f.base }

func (f fakeRegistry) Repo(context.Context, string) (hf.Repo, error) { return f.repo, nil }

// fakeMirror serves the header at every file URL and reports each fetched
// file present, with the checksum it was told to.
type fakeMirror struct {
	mu   sync.Mutex
	srv  *httptest.Server
	got  []mirror.FetchRequest
	sums map[string]string
}

func (m *fakeMirror) Fetch(_ context.Context, r mirror.FetchRequest) (mirror.FetchResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.got = append(m.got, r)
	return mirror.FetchResponse{URL: m.srv.URL + "/" + r.Model + "/" + r.Version + "/" + r.File}, nil
}

func (m *fakeMirror) List(context.Context) (mirror.ModelsResponse, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out mirror.ModelsResponse
	for _, r := range m.got {
		out.Models = append(out.Models, mirror.Status{Model: r.Model, Version: r.Version, File: r.File,
			State: mirror.Present, BytesDone: 1 << 30, BytesTotal: 1 << 30, SHA256: m.sums[r.File]})
	}
	return out, nil
}

type fakeCatalog struct {
	mu    sync.Mutex
	added []catalog.Model
}

func (c *fakeCatalog) Add(_ context.Context, m catalog.Model, _ string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.added = append(c.added, m)
	return nil
}

// fakeFleet has a replica ready at engine once the catalog holds the model.
type fakeFleet struct {
	cat    *fakeCatalog
	engine string
	health []reconciler.ModelHealth
}

func (f *fakeFleet) Replicas() []reconciler.Replica {
	f.cat.mu.Lock()
	defer f.cat.mu.Unlock()
	var out []reconciler.Replica
	for _, m := range f.cat.added {
		out = append(out, reconciler.Replica{ID: m.Name + "-r1", Model: m.Name,
			State: reconciler.Ready, Endpoint: f.engine,
			Shards: []reconciler.Shard{{NodeID: "node01"}}})
	}
	return out
}
func (f *fakeFleet) ModelHealth() []reconciler.ModelHealth { return f.health }
func (f *fakeFleet) EngineKey(string) string               { return "k" }

func setup(t *testing.T, files []hf.File, sums map[string]string) (*Server, *fakeMirror, *fakeCatalog, *fakeFleet) {
	t.Helper()
	header := ggufHeader()
	files0 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write(header)
	}))
	t.Cleanup(files0.Close)
	engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer k" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello!"}}],"usage":{"completion_tokens":3}}`))
	}))
	t.Cleanup(engine.Close)

	m := &fakeMirror{srv: files0, sums: sums}
	c := &fakeCatalog{}
	f := &fakeFleet{cat: c, engine: engine.URL}
	s := New(Deps{
		Registry: fakeRegistry{repo: hf.Repo{ID: "bartowski/Qwen2.5-1.5B-Instruct-GGUF", Revision: "abcdef123456", Files: files}},
		Mirror:   m, Catalog: c, Fleet: f, Poll: 10 * time.Millisecond, ReadyTimeout: 5 * time.Second,
	})
	return s, m, c, f
}

func wait(t *testing.T, s *Server, id string) Job {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if j, _ := s.Get(id); j.State != Running {
			return j
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("the job never finished")
	return Job{}
}

func TestAModelGoesFromTheRegistryToAnAnsweringReplica(t *testing.T) {
	files := []hf.File{
		{Path: "Qwen2.5-1.5B-Instruct-Q8_0.gguf", Bytes: 2 << 30, SHA256: "aa"},
		{Path: "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf", Bytes: 1 << 30, SHA256: "bb"},
		{Path: "mmproj-Q4_K_M.gguf", Bytes: 1 << 20},
	}
	s, m, c, _ := setup(t, files, map[string]string{"Qwen2.5-1.5B-Instruct-Q4_K_M.gguf": "BB"})
	job, err := s.Start(Request{Repo: "bartowski/Qwen2.5-1.5B-Instruct-GGUF", Replicas: 1})
	if err != nil {
		t.Fatal(err)
	}
	j := wait(t, s, job.ID)
	if j.State != Done {
		t.Fatalf("job %s: %s\n%+v", j.State, j.Error, j.Steps)
	}
	if len(m.got) != 1 || m.got[0].File != "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf" || m.got[0].Version != "abcdef12" {
		t.Errorf("fetched %+v; expected the Q4_K_M file at its revision", m.got)
	}
	mod := c.added[0]
	if mod.Name != "qwen2.5-1.5b-instruct" || mod.Profiles[0].MinMemoryGB == 0 || mod.Scaling.Min != 1 {
		t.Errorf("catalog entry %+v", mod)
	}
	if j.Reply != "Hello!" {
		t.Errorf("the smoke test's answer was %q", j.Reply)
	}
	var names []string
	for _, st := range j.Steps {
		names = append(names, st.Name)
	}
	if strings.Join(names, ",") != "registry,mirror,check,catalog,placement,smoke test" {
		t.Errorf("steps %v", names)
	}
}

func TestAQuantTheRepositoryDoesNotHaveIsNamedWithWhatItHas(t *testing.T) {
	s, _, c, _ := setup(t, []hf.File{{Path: "m-Q8_0.gguf", Bytes: 1}}, nil)
	job, _ := s.Start(Request{Repo: "a/b", Replicas: 1})
	j := wait(t, s, job.ID)
	if j.State != Failed || !strings.Contains(j.Error, "m-Q8_0.gguf") {
		t.Errorf("%s: %s", j.State, j.Error)
	}
	if len(c.added) != 0 {
		t.Error("something was added to the catalog")
	}
}

func TestADownloadThatDoesNotMatchTheRegistryIsNotServed(t *testing.T) {
	files := []hf.File{{Path: "m-Q4_K_M.gguf", Bytes: 1 << 30, SHA256: "aaaaaaaaaaaaaaaa"}}
	s, _, c, _ := setup(t, files, map[string]string{"m-Q4_K_M.gguf": "bbbbbbbbbbbbbbbb"})
	job, _ := s.Start(Request{Repo: "a/b", Replicas: 1})
	j := wait(t, s, job.ID)
	if j.State != Failed || !strings.Contains(j.Error, "registry published") {
		t.Errorf("%s: %s", j.State, j.Error)
	}
	if len(c.added) != 0 {
		t.Error("a corrupted download reached the catalog")
	}
}

func TestAParkedModelStopsTheWait(t *testing.T) {
	s, _, _, f := setup(t, []hf.File{{Path: "m-Q4_K_M.gguf", Bytes: 1 << 30}}, nil)
	f.health = []reconciler.ModelHealth{{Model: "qwen2.5-1.5b-instruct", Parked: true, Failures: 5, Reason: "unknown model architecture"}}
	// A fleet where it never becomes ready, only parked.
	s.deps.Fleet = parkedFleet{f}
	job, _ := s.Start(Request{Repo: "a/b", Replicas: 1})
	j := wait(t, s, job.ID)
	if j.State != Failed || !strings.Contains(j.Error, "parked") {
		t.Errorf("%s: %s", j.State, j.Error)
	}
}

type parkedFleet struct{ *fakeFleet }

func (parkedFleet) Replicas() []reconciler.Replica { return nil }

func TestARepositoryThatIsNotOneIsRefused(t *testing.T) {
	s := New(Deps{})
	for _, r := range []string{"", "noslash", "../etc/passwd", "a/b/c"} {
		if _, err := s.Start(Request{Repo: r}); err == nil {
			t.Errorf("%q accepted", r)
		}
	}
}

func TestAVisionModelIsServedWithItsProjectorBesideIt(t *testing.T) {
	files := []hf.File{
		{Path: "gemma-3-4b-it-Q4_K_M.gguf", Bytes: 2 << 30, SHA256: "aa"},
		{Path: "mmproj-gemma-3-4b-it-f16.gguf", Bytes: 800 << 20, SHA256: "bb"},
	}
	s, m, c, _ := setup(t, files, map[string]string{"gemma-3-4b-it-Q4_K_M.gguf": "aa", "mmproj-gemma-3-4b-it-f16.gguf": "bb"})
	job, err := s.Start(Request{Repo: "google/gemma-3-4b-it-GGUF", Quant: "Q4_K_M",
		Projector: "mmproj-gemma-3-4b-it-f16.gguf", Replicas: 1})
	if err != nil {
		t.Fatal(err)
	}
	if j := wait(t, s, job.ID); j.State != Done {
		t.Fatalf("%s %+v", j.Error, j.Steps)
	}
	if len(m.got) != 2 {
		t.Fatalf("fetched %+v", m.got)
	}
	got := c.added[0]
	if !strings.HasSuffix(got.Source, "/") || !got.Capabilities.Vision || len(got.Origin.Files) != 2 {
		t.Errorf("%s vision=%v origin=%+v", got.Source, got.Capabilities.Vision, got.Origin)
	}
	// A projector the repository does not have is named with those it has.
	job, _ = s.Start(Request{Repo: "google/gemma-3-4b-it-GGUF", Quant: "Q4_K_M", Projector: "mmproj-x.gguf", Name: "g2"})
	if j := wait(t, s, job.ID); j.State != Failed || !strings.Contains(j.Error, "mmproj-gemma-3-4b-it-f16.gguf") {
		t.Errorf("%+v", j)
	}
}

func TestOriginalWeightsAreServedOnVLLMWithTheirConfigAndTokenizer(t *testing.T) {
	files := []hf.File{
		{Path: "model-00001-of-00002.safetensors", Bytes: 4 << 30, SHA256: "a1"},
		{Path: "model-00002-of-00002.safetensors", Bytes: 3 << 30, SHA256: "a2"},
		{Path: "model.safetensors.index.json", Bytes: 30000},
		{Path: "config.json", Bytes: 700},
		{Path: "tokenizer.json", Bytes: 7 << 20},
		{Path: "tokenizer_config.json", Bytes: 7000},
		{Path: "README.md", Bytes: 9000},
		{Path: "original/consolidated.safetensors", Bytes: 7 << 30},
	}
	s, m, c, _ := setup(t, files, map[string]string{"model-00001-of-00002.safetensors": "a1", "model-00002-of-00002.safetensors": "a2"})
	// With no fleet to plan against, the GPUs must be named.
	job0, _ := s.Start(Request{Repo: "Qwen/Qwen2.5-7B-Instruct", Format: "safetensors", Name: "x"})
	if j := wait(t, s, job0.ID); j.State != Failed || !strings.Contains(j.Error, "GPU class") {
		t.Errorf("safetensors with no GPU class: %+v", j)
	}
	job, err := s.Start(Request{Repo: "Qwen/Qwen2.5-7B-Instruct", Format: "safetensors", GPUClass: "l40s-48g", Replicas: 1})
	if err != nil {
		t.Fatal(err)
	}
	if j := wait(t, s, job.ID); j.State != Done {
		t.Fatalf("%s %+v", j.Error, j.Steps)
	}
	var got []string
	for _, r := range m.got {
		got = append(got, r.File)
	}
	if len(got) != 6 || slices.Contains(got, "README.md") || slices.Contains(got, "consolidated.safetensors") {
		t.Errorf("fetched %v", got)
	}
	mod := c.added[0]
	p := mod.Profiles[0]
	if !strings.HasSuffix(mod.Source, "/") || p.Engine != "vllm" || p.GPUClass != "l40s-48g" || p.TP != 1 ||
		!slices.Contains(p.Args, "--max-model-len") {
		t.Errorf("%s %+v", mod.Source, p)
	}
}

func TestTheFileThatSuitsTheCardsIsChosenAndRunsOnThem(t *testing.T) {
	files := []hf.File{
		{Path: "m-Q4_K_M.gguf", Bytes: 1 << 30, SHA256: "a"},
		{Path: "m-Q8_0.gguf", Bytes: 2 << 30, SHA256: "b"},
		{Path: "m-F16.gguf", Bytes: 40 << 30, SHA256: "c"},
	}
	s, m, c, _ := setup(t, files, map[string]string{"m-Q4_K_M.gguf": "a", "m-Q8_0.gguf": "b"})
	s.deps.Registry = fakeRegistry{repo: hf.Repo{ID: "o/m-GGUF", Revision: "abcdef123456", Files: files}, base: m.srv.URL}
	s.deps.Pool = func() []modelplan.PoolNode {
		return []modelplan.PoolNode{
			{ID: "g", Class: "6000-24g", GPUMemGB: 24, GPUs: 1, FreeGPUs: 1, HostMemGB: 64, HostFreeGB: 60},
			{ID: "c", HostMemGB: 16, HostFreeGB: 14},
		}
	}
	p, err := s.Plan(context.Background(), PlanRequest{Repo: "o/m-GGUF"})
	if err != nil {
		t.Fatal(err)
	}
	rec, _ := p.Recommended()
	if rec.Artifact != "m-Q8_0.gguf" || rec.CPU || rec.Class != "6000-24g" || p.Artifacts[2].Best == nil || !p.Artifacts[2].Best.CPU {
		t.Fatalf("plan: %+v; f16: %+v", rec, p.Artifacts[2].Best)
	}
	job, err := s.Start(Request{Repo: "o/m-GGUF", Replicas: 1, Fallbacks: true})
	if err != nil {
		t.Fatal(err)
	}
	if j := wait(t, s, job.ID); j.State != Done {
		t.Fatalf("%s %+v", j.Error, j.Steps)
	}
	if m.got[0].File != "m-Q8_0.gguf" {
		t.Errorf("fetched %s", m.got[0].File)
	}
	profs := c.added[0].Profiles
	if len(profs) != 2 || profs[0].GPUClass != "6000-24g" || profs[0].GPUsPerNode != 1 || profs[0].Engine != "llamacpp" ||
		profs[1].Name != "cpu" || profs[1].Priority >= profs[0].Priority || profs[1].MinMemoryGB == 0 {
		t.Errorf("profiles: %+v", profs)
	}
	// A file too big for any card or any node's memory: refused, saying so.
	s.deps.Pool = func() []modelplan.PoolNode { return []modelplan.PoolNode{{ID: "c", HostMemGB: 16, HostFreeGB: 14}} }
	job, _ = s.Start(Request{Repo: "o/m-GGUF", File: "m-F16.gguf", Name: "big"})
	if j := wait(t, s, job.ID); j.State != Failed || !strings.Contains(j.Error, "fits nowhere") {
		t.Errorf("%+v", j)
	}
}

// busyFleet has the model's cards all in use: nothing placed, and placement
// saying so.
type busyFleet struct{ *fakeFleet }

func (busyFleet) Replicas() []reconciler.Replica { return nil }
func (f busyFleet) ModelHealth() []reconciler.ModelHealth {
	return []reconciler.ModelHealth{{Model: "qwen2.5-1.5b-instruct", Unplaceable: true,
		Placement: "cuda01: 0 of 1 6000-24g free"}}
}

func TestAModelThatFitsBusyCardsIsAddedToWaitNotRefused(t *testing.T) {
	files := []hf.File{{Path: "m-Q8_0.gguf", Bytes: 2 << 30, SHA256: "b"}}
	s, m, c, f := setup(t, files, map[string]string{"m-Q8_0.gguf": "b"})
	s.deps.Registry = fakeRegistry{repo: hf.Repo{ID: "o/m-GGUF", Revision: "abcdef123456", Files: files}, base: m.srv.URL}
	s.deps.Pool = func() []modelplan.PoolNode {
		return []modelplan.PoolNode{{ID: "g", Class: "6000-24g", GPUMemGB: 24, GPUs: 1, FreeGPUs: 0, HostMemGB: 8}}
	}
	s.deps.Fleet = busyFleet{f}
	s.deps.UnplacedGrace = 10 * time.Millisecond
	job, err := s.Start(Request{Repo: "o/m-GGUF", Name: "qwen2.5-1.5b-instruct", Replicas: 1})
	if err != nil {
		t.Fatal(err)
	}
	j := wait(t, s, job.ID)
	if j.State != Done || len(c.added) != 1 || c.added[0].Profiles[0].GPUClass != "6000-24g" {
		t.Fatalf("%s %+v", j.Error, j.Steps)
	}
	last := j.Steps[len(j.Steps)-1]
	if last.Name != "placement" || last.State != Skipped || !strings.Contains(last.Detail, "starts when there is room") {
		t.Errorf("%+v", last)
	}
}

func TestOnlyTheLastFinishedJobsAreRemembered(t *testing.T) {
	s := New(Deps{})
	start := time.Now().Add(-time.Hour)
	s.mu.Lock()
	running := &Job{ID: "running", State: Running, Started: start}
	s.keepJob(running)
	for i := range keptJobs + 5 {
		s.keepJob(&Job{ID: fmt.Sprintf("done-%d", i), State: "done", Started: start.Add(time.Duration(i+1) * time.Second)})
	}
	s.mu.Unlock()
	if n := len(s.Jobs()); n != keptJobs+1 {
		t.Errorf("%d jobs remembered, expected %d finished and the running one", n, keptJobs)
	}
	if _, ok := s.Get("running"); !ok {
		t.Error("the running job, the oldest, was forgotten")
	}
	if _, ok := s.Get("done-0"); ok {
		t.Error("the oldest finished job is still remembered")
	}
}