package main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"wxch.org/wxch/internal/modules"
)

type fakeModRunner struct {
	mu       sync.Mutex
	image    map[string]bool
	pulls    []string
	stops    []string
	rmis     []string
	pulling  bool
	running  map[string]bool
	exited   map[string]bool
	startErr string
	logs     string
}

func (f *fakeModRunner) Status() ([]runnerModule, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []runnerModule
	for _, e := range []string{"comfyui", "whisper"} {
		st := ""
		if f.running[e] {
			st = "running"
		}
		if f.exited[e] {
			st = "exited"
		}
		out = append(out, runnerModule{Engine: e, ImagePresent: f.image[e], Pulling: f.pulling, State: st})
	}
	return out, nil
}
func (f *fakeModRunner) Pull(e string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.pulls = append(f.pulls, e)
	f.image[e] = true // the pull "finishes" at once
	return nil
}
func (f *fakeModRunner) Start(e string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.startErr != "" {
		return fmt.Errorf("%s", f.startErr)
	}
	if f.running == nil {
		f.running = map[string]bool{}
	}
	f.running[e] = true
	return nil
}
func (f *fakeModRunner) Stop(e string) error     { f.stops = append(f.stops, e); return nil }
func (f *fakeModRunner) Logs(string, int) string { return f.logs }
func (f *fakeModRunner) Rmi(e string) error      { f.rmis = append(f.rmis, e); return nil }

// fakeHF serves bytes of the given length for every path, with Range support.
func fakeModelServer(t *testing.T) *httptest.Server {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n, _ := strconv.Atoi(r.URL.Query().Get("n"))
		data := []byte(strings.Repeat("x", n))
		http.ServeContent(w, r, "f", time.Time{}, strings.NewReader(string(data)))
	}))
	t.Cleanup(srv.Close)
	return srv
}

func testCatalog(base string) *modules.Catalog {
	mf := func(dest string, n int) modules.ModelFile {
		return modules.ModelFile{Dest: dest, URL: fmt.Sprintf("%s/%s?n=%d", base, dest, n), Bytes: int64(n)}
	}
	return &modules.Catalog{
		Engines: []modules.Engine{
			{ID: "comfyui", Image: "img", Digest: "sha256:aa", Workdir: "/work/ComfyUI", SourceRepo: "https://github.com/x/y", SourceTag: "v1", ModelsRoot: "ComfyUI/models", Port: 1, Cmd: []string{"x"}},
			{ID: "whisper", Image: "img", Digest: "sha256:aa", Workdir: "/work", ModelsRoot: "models", Port: 2, Cmd: []string{"x"}},
		},
		Capabilities: []modules.Capability{
			{ID: "a.one", Engine: "comfyui", Models: []modules.ModelFile{mf("checkpoints/shared.bin", 3000), mf("loras/a.bin", 1000)}},
			{ID: "a.two", Engine: "comfyui", Models: []modules.ModelFile{mf("checkpoints/shared.bin", 3000), mf("vae/b.bin", 500)}},
			{ID: "s.one", Engine: "whisper", Models: []modules.ModelFile{mf("w/model.bin", 2000)}},
		},
	}
}

func moduleSetup(t *testing.T) (nodeConfig, *fakeModRunner, *modules.Catalog) {
	srv := fakeModelServer(t)
	cat := testCatalog(srv.URL)
	old := loadCatalog
	loadCatalog = func() (*modules.Catalog, error) { return cat, nil }
	rn := &fakeModRunner{image: map[string]bool{}}
	modRunnerOverride = rn
	moduleManager = &moduleMgr{last: map[string]*installProgress{}}
	t.Cleanup(func() { loadCatalog = old; modRunnerOverride = nil })
	cfg := nodeConfig{ModulesDir: filepath.Join(t.TempDir(), "modules"), MinFreeGB: 0}
	// the engine code is "already cloned", so no git is needed
	src := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI")
	os.MkdirAll(src, 0755)
	os.WriteFile(filepath.Join(src, "main.py"), []byte("#"), 0644)
	os.WriteFile(filepath.Join(src, "requirements.txt"), []byte("torch\ntorchvision==0.1\ntorchaudio>=2\ntorchsde\nnumpy\r\n"), 0644)
	return cfg, rn, cat
}

func waitIdle(t *testing.T) {
	for i := 0; i < 200; i++ {
		moduleManager.mu.Lock()
		b := moduleManager.busy
		moduleManager.mu.Unlock()
		if b == nil {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the task did not finish")
}

func TestInstallDownloadsEverything(t *testing.T) {
	cfg, rn, _ := moduleSetup(t)
	if err := moduleManager.Install(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	waitIdle(t)
	if l := moduleManager.last["a.one"]; l == nil || l.Err != "" || l.Status != "done" {
		t.Fatalf("task = %+v", l)
	}
	if readModuleState(cfg).Installed["a.one"] == "" {
		t.Fatal("not recorded as installed")
	}
	root := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "models")
	for f, n := range map[string]int64{"checkpoints/shared.bin": 3000, "loras/a.bin": 1000} {
		fi, err := os.Stat(filepath.Join(root, f))
		if err != nil || fi.Size() != n {
			t.Errorf("%s: %v", f, err)
		}
	}
	if len(rn.pulls) != 1 {
		t.Errorf("image pulls = %v", rn.pulls)
	}
	req, _ := os.ReadFile(filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "requirements.lab.txt"))
	if strings.Contains(string(req), "torch\n") || strings.Contains(string(req), "torchvision") || strings.Contains(string(req), "torchaudio") || !strings.Contains(string(req), "torchsde") || !strings.Contains(string(req), "numpy") {
		t.Errorf("requirements.lab.txt = %q", req)
	}
	if err := moduleManager.Install(cfg, "a.one"); err == nil {
		t.Error("installing twice was accepted")
	}
}

func TestInstallResumesAPartFile(t *testing.T) {
	cfg, _, _ := moduleSetup(t)
	root := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "models", "loras")
	os.MkdirAll(root, 0755)
	os.WriteFile(filepath.Join(root, "a.bin.part"), []byte(strings.Repeat("x", 400)), 0644)
	if err := moduleManager.Install(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	waitIdle(t)
	if l := moduleManager.last["a.one"]; l.Err != "" {
		t.Fatal(l.Err)
	}
	if fi, err := os.Stat(filepath.Join(root, "a.bin")); err != nil || fi.Size() != 1000 {
		t.Fatalf("a.bin: %v", err)
	}
}

func TestRemoveKeepsSharedFilesAndCleansUp(t *testing.T) {
	cfg, rn, _ := moduleSetup(t)
	for _, id := range []string{"a.one", "a.two", "s.one"} {
		if err := moduleManager.Install(cfg, id); err != nil {
			t.Fatal(err)
		}
		waitIdle(t)
		if l := moduleManager.last[id]; l.Err != "" {
			t.Fatal(l.Err)
		}
	}
	root := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "models")
	if err := moduleManager.Remove(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(root, "checkpoints", "shared.bin")); err != nil {
		t.Error("a file another capability needs was deleted")
	}
	if _, err := os.Stat(filepath.Join(root, "loras", "a.bin")); err == nil {
		t.Error("the capability's own file stayed")
	}
	if _, err := os.Stat(filepath.Join(root, "loras")); err == nil {
		t.Error("the empty folder stayed")
	}
	if len(rn.rmis) != 0 {
		t.Errorf("image removed while engines still use it: %v", rn.rmis)
	}
	moduleManager.Remove(cfg, "a.two")
	if _, err := os.Stat(filepath.Join(cfg.ModulesDir, "comfyui")); err == nil {
		t.Error("the engine folder stayed after its last capability left")
	}
	if len(rn.rmis) != 0 {
		t.Errorf("the image is shared with whisper, which is still installed: %v", rn.rmis)
	}
	moduleManager.Remove(cfg, "s.one")
	if len(rn.rmis) != 1 {
		t.Errorf("image should go when no engine needs it: %v", rn.rmis)
	}
	if len(readModuleState(cfg).Installed) != 0 {
		t.Error("state not empty")
	}
}

func TestRemoveLeavesAFolderThatIsNotOurs(t *testing.T) {
	cfg, _, _ := moduleSetup(t)
	moduleManager.Install(cfg, "s.one")
	waitIdle(t)
	dir := filepath.Join(cfg.ModulesDir, "whisper")
	os.Remove(filepath.Join(dir, moduleMarker))
	os.WriteFile(filepath.Join(dir, "mine.txt"), []byte("x"), 0644)
	moduleManager.Remove(cfg, "s.one")
	if _, err := os.Stat(filepath.Join(dir, "mine.txt")); err != nil {
		t.Error("a folder without our marker was deleted")
	}
}

func TestOneTaskAtATimeAndUnknownIDs(t *testing.T) {
	cfg, _, _ := moduleSetup(t)
	if err := moduleManager.Install(cfg, "nope"); err == nil {
		t.Error("unknown capability accepted")
	}
	if err := moduleManager.Remove(cfg, "a.one"); err == nil {
		t.Error("removing something that is not installed was accepted")
	}
	moduleManager.mu.Lock()
	moduleManager.busy = &installProgress{ID: "x", Op: "install"}
	moduleManager.mu.Unlock()
	if err := moduleManager.Install(cfg, "a.one"); err == nil {
		t.Error("a second task was accepted")
	}
	moduleManager.mu.Lock()
	moduleManager.busy = nil
	moduleManager.mu.Unlock()
}

func TestOverviewListsEverything(t *testing.T) {
	cfg, _, cat := moduleSetup(t)
	moduleManager.Install(cfg, "s.one")
	waitIdle(t)
	o := moduleManager.Overview(cfg)
	if len(o.Capabilities) != len(cat.Capabilities) || len(o.Engines) != 2 {
		t.Fatalf("overview = %+v", o)
	}
	for _, c := range o.Capabilities {
		if c.ID == "s.one" && !c.Installed {
			t.Error("s.one not shown as installed")
		}
	}
}

func TestRemoveNeverDeletesModelsTheAdministratorAdded(t *testing.T) {
	cfg, _, _ := moduleSetup(t)
	if err := moduleManager.Install(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	waitIdle(t)
	models := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "models")
	os.MkdirAll(filepath.Join(models, "diffusion_models"), 0755)
	os.MkdirAll(filepath.Join(models, "configs"), 0755)
	os.WriteFile(filepath.Join(models, "diffusion_models", "my-own.safetensors"), []byte("MINE"), 0644)
	os.WriteFile(filepath.Join(models, "configs", "sample.yaml"), []byte("x"), 0644)
	os.WriteFile(filepath.Join(models, "diffusion_models", "put_diffusion_model_files_here"), nil, 0644)
	if err := moduleManager.Remove(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	if b, err := os.ReadFile(filepath.Join(models, "diffusion_models", "my-own.safetensors")); err != nil || string(b) != "MINE" {
		t.Fatal("a model the administrator added was deleted")
	}
	if _, err := os.Stat(filepath.Join(models, "loras", "a.bin")); err == nil {
		t.Error("the module's own file stayed")
	}
	// with nothing of the administrator's, the folder goes as before (placeholder and config files do not count)
	os.Remove(filepath.Join(models, "diffusion_models", "my-own.safetensors"))
	moduleManager.Install(cfg, "a.one")
	waitIdle(t)
	if err := moduleManager.Remove(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(cfg.ModulesDir, "comfyui")); err == nil {
		t.Error("the engine folder stayed although it held nothing of the administrator's")
	}
}

func TestExtraFilesAreListedAndOnlyThoseCanBeDeleted(t *testing.T) {
	cfg, _, cat := moduleSetup(t)
	if err := moduleManager.Install(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	waitIdle(t)
	models := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "models")
	os.MkdirAll(filepath.Join(models, "checkpoints"), 0755)
	os.MkdirAll(filepath.Join(models, "text_encoders"), 0755)
	os.MkdirAll(filepath.Join(models, "configs"), 0755)
	os.WriteFile(filepath.Join(models, "checkpoints", "ltx.safetensors"), []byte(strings.Repeat("x", 1234)), 0644)
	os.WriteFile(filepath.Join(models, "text_encoders", "gemma.safetensors"), []byte("g"), 0644)
	os.WriteFile(filepath.Join(models, "checkpoints", "half.safetensors.part"), []byte("p"), 0644)
	os.WriteFile(filepath.Join(models, "checkpoints", "put_checkpoints_here"), nil, 0644)
	os.WriteFile(filepath.Join(models, "configs", "v1.yaml"), []byte("c"), 0644)

	got := listExtraFiles(cfg, cat)
	names := map[string]int64{}
	for _, f := range got {
		names[f.Engine+"/"+f.Path] = f.Size
	}
	if len(got) != 2 || names["comfyui/checkpoints/ltx.safetensors"] != 1234 || names["comfyui/text_encoders/gemma.safetensors"] != 1 {
		t.Fatalf("only the added files may be listed (not module files, placeholders, configs or parts): %+v", got)
	}
	if o := moduleManager.Overview(cfg); len(o.Extra) != 2 {
		t.Errorf("the overview must carry them: %+v", o.Extra)
	}

	// what is not an added file cannot be deleted through this path
	for _, bad := range []struct{ eng, path string }{
		{"comfyui", "checkpoints/shared.bin"}, // a module's file
		{"comfyui", "checkpoints/put_checkpoints_here"},
		{"comfyui", "configs/v1.yaml"},
		{"comfyui", "../requirements.txt"},
		{"comfyui", "/etc/passwd"},
		{"comfyui", "checkpoints"}, // a folder
		{"comfyui", "checkpoints/nothing.safetensors"},
		{"nope", "checkpoints/ltx.safetensors"},
	} {
		if err := moduleManager.DeleteExtra(cfg, bad.eng, bad.path); err == nil {
			t.Errorf("%s/%s was accepted", bad.eng, bad.path)
		}
	}
	if _, err := os.Stat(filepath.Join(models, "checkpoints", "shared.bin")); err != nil {
		t.Fatal("a module file was deleted")
	}

	// not while a job runs
	jobs.Lock()
	jobs.running["x"] = &jobRun{info: jobStatusInfo{ID: "x", Capability: "image.generate"}}
	jobs.Unlock()
	if err := moduleManager.DeleteExtra(cfg, "comfyui", "checkpoints/ltx.safetensors"); err == nil {
		t.Error("a file was deleted while a job runs")
	}
	jobs.Lock()
	delete(jobs.running, "x")
	jobs.Unlock()

	if err := moduleManager.DeleteExtra(cfg, "comfyui", "checkpoints/ltx.safetensors"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(models, "checkpoints", "ltx.safetensors")); err == nil {
		t.Error("the file is still there")
	}
	if err := moduleManager.DeleteExtra(cfg, "comfyui", "text_encoders/gemma.safetensors"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(models, "text_encoders")); err == nil {
		t.Error("the empty folder stayed")
	}
	if _, err := os.Stat(filepath.Join(models, "checkpoints", "shared.bin")); err != nil {
		t.Error("deleting added files touched a module file")
	}
	if len(listExtraFiles(cfg, cat)) != 0 {
		t.Error("nothing should be listed any more")
	}
}

func TestStateChangesFromSeveralPlacesDoNotUndoEachOther(t *testing.T) {
	cfg := nodeConfig{ModulesDir: t.TempDir()}
	writeModuleState(cfg, moduleState{Installed: map[string]string{}})
	var wg sync.WaitGroup
	for i := 0; i < 40; i++ {
		wg.Add(2)
		go func(i int) {
			defer wg.Done()
			updateModuleState(cfg, func(st *moduleState) { st.Installed[fmt.Sprintf("cap%d", i)] = "x" })
		}(i)
		go func(i int) {
			defer wg.Done()
			setPinned(cfg, fmt.Sprintf("eng%d", i), true)
		}(i)
	}
	wg.Wait()
	st := readModuleState(cfg)
	if len(st.Installed) != 40 || len(st.Pinned) != 40 {
		t.Fatalf("an update was lost: %d installed, %d pinned", len(st.Installed), len(st.Pinned))
	}
	left, _ := filepath.Glob(filepath.Join(cfg.ModulesDir, "state-*.tmp"))
	if len(left) != 0 {
		t.Errorf("temporary files were left behind: %v", left)
	}
}

func TestADamagedStateFileIsNeverOverwrittenByAGuess(t *testing.T) {
	cfg := nodeConfig{ModulesDir: t.TempDir()}
	writeModuleState(cfg, moduleState{Installed: map[string]string{"image.generate": "x"}})
	writeModuleState(cfg, moduleState{Installed: map[string]string{"image.generate": "x", "video.generate": "y"}})
	if b, err := os.ReadFile(filepath.Join(cfg.ModulesDir, "state.json.bak")); err != nil || !strings.Contains(string(b), "image.generate") || strings.Contains(string(b), "video.generate") {
		t.Fatalf("the previous good file must be kept: %v %q", err, b)
	}
	os.WriteFile(filepath.Join(cfg.ModulesDir, "state.json"), []byte("{ this is not json"), 0644)
	if err := updateModuleState(cfg, func(st *moduleState) { st.Installed["x"] = "y" }); err == nil {
		t.Fatal("a damaged file was overwritten")
	}
	if b, _ := os.ReadFile(filepath.Join(cfg.ModulesDir, "state.json")); string(b) != "{ this is not json" {
		t.Error("the damaged file was changed")
	}
	if len(readModuleState(cfg).Installed) != 0 {
		t.Error("looking at a damaged file shows an empty state, not a guess")
	}
}

func TestARemovalThatFailsLeavesTheStateSayingInstalled(t *testing.T) {
	cfg, rn, cat := moduleSetup(t)
	if err := moduleManager.Install(cfg, "a.one"); err != nil {
		t.Fatal(err)
	}
	waitIdle(t)
	_ = rn
	// make the model file undeletable: replace it with a non-empty folder
	f := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI", "models", "loras", "a.bin")
	os.Remove(f)
	os.MkdirAll(filepath.Join(f, "inside"), 0755)
	os.WriteFile(filepath.Join(f, "inside", "x"), []byte("x"), 0644)
	if err := moduleManager.Remove(cfg, "a.one"); err == nil {
		t.Fatal("the removal should have failed")
	}
	if _, ok := readModuleState(cfg).Installed["a.one"]; !ok {
		t.Fatal("the state said removed although the cleanup failed")
	}
	_ = cat
	os.RemoveAll(f)
	if err := moduleManager.Remove(cfg, "a.one"); err != nil {
		t.Fatalf("a retry must work: %v", err)
	}
	if _, ok := readModuleState(cfg).Installed["a.one"]; ok {
		t.Error("removed, but the state still lists it")
	}
}