package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

func testArchive(files map[string]string) []byte {
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	tw.WriteHeader(&tar.Header{Name: "ComfyUI-0.38.2/", Typeflag: tar.TypeDir, Mode: 0755})
	for name, data := range files {
		tw.WriteHeader(&tar.Header{Name: name, Typeflag: tar.TypeReg, Mode: 0644, Size: int64(len(data))})
		tw.Write([]byte(data))
	}
	tw.Close()
	gz.Close()
	return buf.Bytes()
}

func archiveFixture(t *testing.T, archive []byte, status int) (nodeConfig, modules.Engine) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if status != 200 {
			http.Error(w, "no", status)
			return
		}
		w.Write(archive)
	}))
	t.Cleanup(srv.Close)
	oldURL, oldGit := codeArchiveURL, gitLookPath
	codeArchiveURL = func(modules.Engine) string { return srv.URL + "/x.tar.gz" }
	gitLookPath = func(string) (string, error) { return "", fmt.Errorf("not installed") }
	t.Cleanup(func() { codeArchiveURL, gitLookPath = oldURL, oldGit })
	cfg := nodeConfig{ModulesDir: t.TempDir()}
	return cfg, modules.Engine{ID: "comfyui", Workdir: "/work/ComfyUI", SourceRepo: "https://github.com/comfyanonymous/ComfyUI", SourceTag: "v0.38.2", ModelsRoot: "ComfyUI/models"}
}

func TestEngineCodeComesFromTheArchiveWhenGitIsMissing(t *testing.T) {
	ar := testArchive(map[string]string{
		"ComfyUI-0.38.2/main.py":                "print('hi')",
		"ComfyUI-0.38.2/requirements.txt":       "torch\ntorchsde\nnumpy\n",
		"ComfyUI-0.38.2/models/put_models_here": "",
		"ComfyUI-0.38.2/../evil.txt":            "bad",
		"ComfyUI-0.38.2/comfy/../../escape.txt": "bad",
	})
	cfg, e := archiveFixture(t, ar, 200)
	if err := prepareEngine(cfg, e); err != nil {
		t.Fatal(err)
	}
	src := filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI")
	if b, err := os.ReadFile(filepath.Join(src, "main.py")); err != nil || string(b) != "print('hi')" {
		t.Fatalf("main.py: %q %v", b, err)
	}
	req, _ := os.ReadFile(filepath.Join(src, "requirements.lab.txt"))
	if strings.Contains(string(req), "torch\n") || !strings.Contains(string(req), "torchsde") || !strings.Contains(string(req), "numpy") {
		t.Errorf("requirements.lab.txt = %q", req)
	}
	for _, bad := range []string{filepath.Join(cfg.ModulesDir, "comfyui", "evil.txt"), filepath.Join(cfg.ModulesDir, "escape.txt"), filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI.unpack")} {
		if _, err := os.Stat(bad); err == nil {
			t.Errorf("%s exists", bad)
		}
	}
	if _, err := os.Stat(filepath.Join(cfg.ModulesDir, "comfyui", moduleMarker)); err != nil {
		t.Error("the engine folder is not marked as ours")
	}
	// running it again with the code in place does not fetch anything
	codeArchiveURL = func(modules.Engine) string { return "http://127.0.0.1:1/never" }
	if err := prepareEngine(cfg, e); err != nil {
		t.Fatalf("a second run: %v", err)
	}
}

func TestArchiveProblemsAreReportedAndLeaveNothingHalfDone(t *testing.T) {
	cfg, e := archiveFixture(t, nil, 404)
	if err := prepareEngine(cfg, e); err == nil || !strings.Contains(err.Error(), "git is not installed") {
		t.Errorf("a missing archive: %v", err)
	}
	cfg, e = archiveFixture(t, testArchive(map[string]string{"ComfyUI-0.38.2/readme.txt": "no engine here"}), 200)
	if err := prepareEngine(cfg, e); err == nil {
		t.Error("an archive without main.py was accepted")
	}
	if _, err := os.Stat(filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI")); err == nil {
		t.Error("a half-unpacked folder was left in place")
	}
	cfg, e = archiveFixture(t, []byte("this is not gzip"), 200)
	if err := prepareEngine(cfg, e); err == nil {
		t.Error("garbage was accepted")
	}
}