package main

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

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

// TestHelperEngine is not a test: started by the tests below it plays an engine (listens on the port it is given, writes
// its environment to the log, ends on a signal).
func TestHelperEngine(t *testing.T) {
	var port string
	for i, a := range os.Args {
		if a == "wxch-helper-engine" && i+1 < len(os.Args) {
			port = os.Args[i+1]
		}
	}
	if port == "" {
		t.Skip("only run as an engine by the host tests")
	}
	ln, err := net.Listen("tcp", "127.0.0.1:"+port)
	if err != nil {
		fmt.Println("cannot listen:", err)
		os.Exit(3)
	}
	for _, e := range os.Environ() {
		fmt.Println("ENV", e)
	}
	fmt.Println("serving")
	for {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		c.Close()
	}
}

func freePort(t *testing.T) int {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}

func hostTestRunner(t *testing.T) (hostRunner, modules.Engine) {
	t.Helper()
	dir := t.TempDir()
	cat, err := modules.Load()
	if err != nil {
		t.Fatal(err)
	}
	e, _ := cat.Engine("comfyui")
	h := hostRunner{cfg: nodeConfig{ModulesDir: dir}, cat: cat}
	host.Lock()
	host.procs, host.installs = map[string]*hostProc{}, map[string]*hostInstall{}
	host.Unlock()
	t.Cleanup(func() { h.Stop("comfyui"); h.Stop("whisper") })
	return h, e
}

func markReady(t *testing.T, h hostRunner, engine string) {
	t.Helper()
	p := h.paths(engine)
	os.MkdirAll(p.venv, 0755)
	os.MkdirAll(p.dir, 0755)
	os.WriteFile(filepath.Join(p.dir, moduleMarker), nil, 0644)
	os.WriteFile(filepath.Join(p.venv, hostReady), []byte("torch test\n"), 0644)
}

func withHelperEngine(t *testing.T, port int) {
	old := engineArgv
	engineArgv = func(cat *modules.Catalog, e modules.Engine, p hostPaths) ([]string, string) {
		return []string{os.Args[0], "-test.run=TestHelperEngine", "--", "wxch-helper-engine", strconv.Itoa(port)}, p.dir
	}
	t.Cleanup(func() { engineArgv = old })
	// the helper is not a venv python: nothing to adopt or kill by name
}

func waitFor(t *testing.T, what string, f func() bool) {
	t.Helper()
	for i := 0; i < 100; i++ {
		if f() {
			return
		}
		time.Sleep(100 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func TestEngineEnvHoldsNothingFromTheAgent(t *testing.T) {
	t.Setenv("WXCH_TOKEN", "secret-token")
	t.Setenv("SSH_AUTH_SOCK", "/tmp/agent")
	h, _ := hostTestRunner(t)
	env := strings.Join(engineEnv(h.paths("comfyui")), "\n")
	for _, bad := range []string{"WXCH_TOKEN", "secret-token", "SSH_AUTH_SOCK"} {
		if strings.Contains(env, bad) {
			t.Fatalf("the engine environment leaks %s:\n%s", bad, env)
		}
	}
	for _, want := range []string{"HOME=" + h.paths("comfyui").home, "HF_HUB_OFFLINE=1", "PYTHONNOUSERSITE=1"} {
		if !strings.Contains(env, want) {
			t.Errorf("missing %s in\n%s", want, env)
		}
	}
	if ie := strings.Join(installEnv(h.paths("comfyui")), "\n"); strings.Contains(ie, "secret-token") || !strings.Contains(ie, "UV_CACHE_DIR="+filepath.Join(h.paths("comfyui").hostDir, "cache")) {
		t.Errorf("install environment wrong:\n%s", ie)
	}
}

func TestEngineArgvListensOnLoopbackOnTheHostPort(t *testing.T) {
	h, e := hostTestRunner(t)
	argv, dir := engineArgv(h.cat, e, h.paths("comfyui"))
	line := strings.Join(argv, " ")
	if !strings.Contains(line, "--listen 127.0.0.1") || !strings.Contains(line, "--port "+strconv.Itoa(h.cat.HostPort("comfyui"))) || !strings.Contains(line, "--disable-api-nodes") {
		t.Errorf("comfyui: %s", line)
	}
	if filepath.Base(dir) != "ComfyUI" || !strings.HasSuffix(argv[0], filepath.Join("venv", "bin", "python")) {
		t.Errorf("comfyui runs in %s with %s", dir, argv[0])
	}
	we, _ := h.cat.Engine("whisper")
	argv, _ = engineArgv(h.cat, we, h.paths("whisper"))
	line = strings.Join(argv, " ")
	if !strings.Contains(line, "--host 127.0.0.1") || !strings.Contains(line, "--port "+strconv.Itoa(h.cat.HostPort("whisper"))) || !strings.Contains(line, "whisper-large-v3-turbo") {
		t.Errorf("whisper: %s", line)
	}
}

func TestHostEngineStartsAnswersAndStops(t *testing.T) {
	t.Setenv("WXCH_TOKEN", "secret-token")
	h, _ := hostTestRunner(t)
	markReady(t, h, "comfyui")
	port := freePort(t)
	withHelperEngine(t, port)
	if err := h.Start("comfyui"); err != nil {
		t.Fatal(err)
	}
	waitFor(t, "the engine to listen", func() bool {
		c, err := net.DialTimeout("tcp", "127.0.0.1:"+strconv.Itoa(port), 200*time.Millisecond)
		if err == nil {
			c.Close()
		}
		return err == nil
	})
	st, _ := h.Status()
	if st[0].Engine != "comfyui" || st[0].State != "running" || !st[0].ImagePresent {
		t.Fatalf("status: %+v", st[0])
	}
	if err := h.Start("comfyui"); err != nil { // already running: nothing happens
		t.Fatal(err)
	}
	log := h.Logs("comfyui", 200)
	if !strings.Contains(log, "serving") {
		t.Errorf("log: %s", log)
	}
	if strings.Contains(log, "WXCH_TOKEN") || strings.Contains(log, "secret-token") {
		t.Errorf("the engine saw the agent's token:\n%s", log)
	}
	if !strings.Contains(log, "ENV HOME="+h.paths("comfyui").home) {
		t.Errorf("the engine does not run with its own HOME:\n%s", log)
	}
	h.Stop("comfyui")
	waitFor(t, "the engine to stop", func() bool { st, _ := h.Status(); return st[0].State == "exited" })
	if _, err := os.Stat(h.paths("comfyui").pid); err == nil {
		t.Error("the pid file stays after a stop")
	}
	if err := h.Start("comfyui"); err != nil { // and starts again
		t.Fatal(err)
	}
	waitFor(t, "the engine to run again", func() bool { st, _ := h.Status(); return st[0].State == "running" })
}

func TestHostEngineRefusesAPortThatIsTaken(t *testing.T) {
	h, _ := hostTestRunner(t)
	markReady(t, h, "comfyui")
	ln, err := net.Listen("tcp", "127.0.0.1:"+strconv.Itoa(h.cat.HostPort("comfyui")))
	if err != nil {
		t.Skip("the engine port is in use on this machine")
	}
	defer ln.Close()
	if err := h.Start("comfyui"); err == nil || !strings.Contains(err.Error(), "already used") {
		t.Fatalf("got %v", err)
	}
}

func TestHostEngineNotReadyInstallsInsteadOfStarting(t *testing.T) {
	h, e := hostTestRunner(t)
	codeFolder(t, h, e) // the engine code is "already there": the install must not try to download it
	oldUV, oldRun := hostEnsureUV, runLoggedFn
	hostEnsureUV = func(context.Context, string) (string, error) { return "", fmt.Errorf("offline") }
	defer func() { hostEnsureUV, runLoggedFn = oldUV, oldRun }()
	err := h.Start("comfyui")
	if err == nil || !strings.Contains(err.Error(), "still being installed") {
		t.Fatalf("got %v", err)
	}
	waitFor(t, "the install to fail", func() bool { st, _ := h.Status(); return !st[0].Pulling && st[0].PullErr != "" })
	st, _ := h.Status()
	if st[0].ImagePresent || !strings.Contains(st[0].PullErr, "offline") {
		t.Errorf("status: %+v", st[0])
	}
}

// codeFolder puts a fake engine source where prepareEngine finds it, so that no download is tried.
func codeFolder(t *testing.T, h hostRunner, e modules.Engine) {
	t.Helper()
	src := filepath.Join(h.paths(e.ID).dir, filepath.Base(e.Workdir))
	os.MkdirAll(src, 0755)
	os.WriteFile(filepath.Join(src, "main.py"), []byte("print()\n"), 0644)
	os.WriteFile(filepath.Join(src, "requirements.txt"), []byte("torch==2.0\ntorchaudio\nnumpy\nsafetensors>=0.4\n"), 0644)
}

type recorded struct {
	mu   sync.Mutex
	cmds []string
}

func (r *recorded) add(s string) { r.mu.Lock(); r.cmds = append(r.cmds, s); r.mu.Unlock() }

func fakeCommands(t *testing.T, failOn string, r *recorded) {
	oldUV, oldRun := hostEnsureUV, runLoggedFn
	hostEnsureUV = func(_ context.Context, dir string) (string, error) { return filepath.Join(dir, "uv"), nil }
	runLoggedFn = func(_ context.Context, _ io.Writer, dir string, env []string, name string, argv ...string) (string, error) {
		line := filepath.Base(name) + " " + strings.Join(argv, " ")
		r.add(line)
		if failOn != "" && strings.Contains(line, failOn) {
			return "", fmt.Errorf("exit status 1: boom")
		}
		switch {
		case strings.Contains(line, "pip freeze"):
			return "numpy==2.1.0\ntorch==2.14.1+cu130\ntorchvision==0.29.1+cu130\ntorchaudio==2.11.0+cu130\nnvidia-cublas==13.0\n", nil
		case strings.Contains(line, "-c import torch"):
			return "torch 2.14.1 on NVIDIA GB10 ['sm_121']\n", nil
		}
		if strings.HasPrefix(line, "uv venv") {
			os.MkdirAll(filepath.Join(argv[len(argv)-1], "bin"), 0755)
			os.WriteFile(hostPython(argv[len(argv)-1]), nil, 0755)
		}
		return "", nil
	}
	t.Cleanup(func() { hostEnsureUV, runLoggedFn = oldUV, oldRun })
}

func TestProvisionRunsTheStepsInOrderAndPinsTorch(t *testing.T) {
	h, e := hostTestRunner(t)
	codeFolder(t, h, e)
	rec := &recorded{}
	fakeCommands(t, "", rec)
	var steps []string
	if err := h.provision(context.Background(), e, func(s string) { steps = append(steps, s) }); err != nil {
		t.Fatal(err)
	}
	p := h.paths("comfyui")
	want := []string{
		"uv venv --python 3.12 " + p.venv,
		"uv pip install --python " + hostPython(p.venv) + " --torch-backend cu130 torch torchvision torchaudio",
		"uv pip freeze --python " + hostPython(p.venv),
		"uv pip install --python " + hostPython(p.venv) + " -r " + filepath.Join(p.dir, "ComfyUI", "requirements.lab.txt") + " -c " + filepath.Join(p.dir, "constraints.txt"),
		"python -c import torch",
	}
	if len(rec.cmds) != len(want) {
		t.Fatalf("commands:\n%s", strings.Join(rec.cmds, "\n"))
	}
	for i, w := range want {
		if !strings.HasPrefix(rec.cmds[i], w) {
			t.Errorf("command %d:\n got %s\nwant %s", i, rec.cmds[i], w)
		}
	}
	cons, _ := os.ReadFile(filepath.Join(p.dir, "constraints.txt"))
	if strings.TrimSpace(string(cons)) != "torch==2.14.1+cu130\ntorchvision==0.29.1+cu130\ntorchaudio==2.11.0+cu130" {
		t.Errorf("constraints:\n%s", cons)
	}
	// the engine's own requirements do not name torch (the lab file is made from the engine's list)
	lab, _ := os.ReadFile(filepath.Join(p.dir, "ComfyUI", "requirements.lab.txt"))
	if strings.Contains(string(lab), "torch") || !strings.Contains(string(lab), "numpy") {
		t.Errorf("requirements.lab.txt:\n%s", lab)
	}
	if !h.ready("comfyui") {
		t.Error("not marked ready")
	}
	if b, _ := os.ReadFile(filepath.Join(p.venv, hostReady)); !strings.Contains(string(b), "GB10") {
		t.Errorf("the ready marker should say what was verified: %q", b)
	}
	if len(steps) < 5 {
		t.Errorf("steps shown: %v", steps)
	}
}

func TestProvisionStopsAtTheFirstFailureAndNeverMarksReady(t *testing.T) {
	for _, tc := range []struct{ fail, msg string }{
		{"--torch-backend", "pytorch"},
		{"-r ", "packages"},
		{"-c import torch", "GPU check failed"},
	} {
		h, e := hostTestRunner(t)
		codeFolder(t, h, e)
		rec := &recorded{}
		fakeCommands(t, tc.fail, rec)
		err := h.provision(context.Background(), e, func(string) {})
		if err == nil || !strings.Contains(err.Error(), tc.msg) {
			t.Errorf("%s: got %v", tc.fail, err)
		}
		if h.ready("comfyui") {
			t.Errorf("%s: marked ready after a failure", tc.fail)
		}
	}
}

func TestProvisionAfterAFailureContinuesWithoutRedoingPython(t *testing.T) {
	h, e := hostTestRunner(t)
	codeFolder(t, h, e)
	rec := &recorded{}
	fakeCommands(t, "--torch-backend", rec)
	h.provision(context.Background(), e, func(string) {})
	rec2 := &recorded{}
	fakeCommands(t, "", rec2)
	if err := h.provision(context.Background(), e, func(string) {}); err != nil {
		t.Fatal(err)
	}
	for _, c := range rec2.cmds {
		if strings.HasPrefix(c, "uv venv") {
			t.Errorf("the environment was created again: %s", c)
		}
	}
	if !h.ready("comfyui") {
		t.Error("not ready")
	}
}

func TestEnsureRuntimeWaitsForTheInstall(t *testing.T) {
	h, e := hostTestRunner(t)
	codeFolder(t, h, e)
	rec := &recorded{}
	fakeCommands(t, "", rec)
	var reports []string
	if err := h.EnsureRuntime(context.Background(), e, func(s string, _ float64) { reports = append(reports, s) }); err != nil {
		t.Fatal(err)
	}
	if !h.ready("comfyui") {
		t.Fatal("not ready after waiting")
	}
	if err := h.EnsureRuntime(context.Background(), e, nil); err != nil { // ready: no wait, no report
		t.Fatal(err)
	}

	h2, e2 := hostTestRunner(t)
	codeFolder(t, h2, e2)
	fakeCommands(t, "--torch-backend", &recorded{})
	err := h2.EnsureRuntime(context.Background(), e2, func(string, float64) {})
	if err == nil || !strings.Contains(err.Error(), "pytorch") {
		t.Fatalf("a failed install must be reported to the job, got %v", err)
	}
}

func TestRmiRemovesOnlyOurEnvironment(t *testing.T) {
	h, _ := hostTestRunner(t)
	markReady(t, h, "comfyui")
	p := h.paths("comfyui")
	os.MkdirAll(filepath.Join(p.hostDir, "cache"), 0755)
	os.MkdirAll(filepath.Join(p.dir, "ComfyUI", "models"), 0755)
	os.WriteFile(filepath.Join(p.dir, "ComfyUI", "models", "keep.safetensors"), []byte("x"), 0644)
	if err := h.Rmi("comfyui"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(p.venv); err == nil {
		t.Error("the environment stays")
	}
	if _, err := os.Stat(p.hostDir); err == nil {
		t.Error("the shared uv folder stays although no engine has an environment")
	}
	if _, err := os.Stat(filepath.Join(p.dir, "ComfyUI", "models", "keep.safetensors")); err != nil {
		t.Error("model files must not be touched")
	}
	// a folder without our marker is not ours
	h2, _ := hostTestRunner(t)
	q := h2.paths("whisper")
	os.MkdirAll(q.venv, 0755)
	if err := h2.Rmi("whisper"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(q.venv); err != nil {
		t.Error("a folder that is not ours was changed")
	}
}

func TestLogsTailAndTheInstallNotice(t *testing.T) {
	h, _ := hostTestRunner(t)
	if got := h.Logs("comfyui", 5); got != "(no log yet)" {
		t.Errorf("no log: %q", got)
	}
	p := h.paths("comfyui")
	os.MkdirAll(p.dir, 0755)
	var b strings.Builder
	for i := 0; i < 100; i++ {
		fmt.Fprintf(&b, "line %d\n", i)
	}
	os.WriteFile(p.log, []byte(b.String()), 0644)
	got := h.Logs("comfyui", 3)
	if got != "line 97\nline 98\nline 99" {
		t.Errorf("tail: %q", got)
	}
}

func TestEngineModeDefaultsToHostAndCanBeChosen(t *testing.T) {
	dir := t.TempDir()
	cfg := nodeConfig{ModulesDir: dir}
	if engineMode(cfg) != "host" {
		t.Fatal("the default is the engines running directly on the machine")
	}
	cfg.EngineMode = "docker"
	if engineMode(cfg) != "docker" {
		t.Error("node config")
	}
	writeModuleState(cfg, moduleState{Installed: map[string]string{}, Mode: "host"})
	if engineMode(cfg) != "host" {
		t.Error("state.json wins over the node config")
	}
	cfg.EngineMode = "nonsense"
	writeModuleState(cfg, moduleState{Installed: map[string]string{}, Mode: "also nonsense"})
	if engineMode(cfg) != "host" {
		t.Error("an unknown mode is ignored")
	}
}

func uvArchive(t *testing.T, content string) ([]byte, string) {
	t.Helper()
	var buf strings.Builder
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	tw.WriteHeader(&tar.Header{Name: "uv-x/", Typeflag: tar.TypeDir, Mode: 0755})
	tw.WriteHeader(&tar.Header{Name: "uv-x/uvx", Typeflag: tar.TypeReg, Mode: 0755, Size: 3})
	tw.Write([]byte("uvx"))
	tw.WriteHeader(&tar.Header{Name: "uv-x/uv", Typeflag: tar.TypeReg, Mode: 0755, Size: int64(len(content))})
	tw.Write([]byte(content))
	tw.Close()
	gz.Close()
	return []byte(buf.String()), ""
}

func TestUVIsCheckedAgainstThePinnedChecksum(t *testing.T) {
	data, _ := uvArchive(t, "#!/bin/sh\necho uv 0.12.23\n")
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write(data) }))
	defer srv.Close()
	oldURL, oldSums := uvURL, uvSums
	uvURL = func(string) string { return srv.URL }
	defer func() { uvURL, uvSums = oldURL, oldSums }()

	uvSums = map[string]string{runtime.GOARCH: strings.Repeat("0", 64)}
	dir := t.TempDir()
	if _, err := ensureUV(context.Background(), dir); err == nil || !strings.Contains(err.Error(), "checksum") {
		t.Fatalf("a download that does not match the pinned checksum must be refused, got %v", err)
	}
	if _, err := os.Stat(filepath.Join(dir, "uv")); err == nil {
		t.Fatal("a refused download must not leave a program behind")
	}

	uvSums = map[string]string{runtime.GOARCH: sumHex(data)}
	got, err := ensureUV(context.Background(), dir)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := os.ReadFile(got)
	if string(b) != "#!/bin/sh\necho uv 0.12.23\n" {
		t.Errorf("unpacked: %q", b) // the file called uv, not uvx
	}
}

func TestPinnedUVChecksumsAreForTheBuildsWeDownload(t *testing.T) {
	for arch, triple := range uvTriples {
		if len(uvSums[arch]) != 64 {
			t.Errorf("%s has no pinned checksum", arch)
		}
		if u := uvURL(arch); !strings.Contains(u, "/"+uvVersion+"/") || !strings.Contains(u, triple) {
			t.Errorf("url: %s", u)
		}
	}
}

func sumHex(b []byte) string { s := sha256.Sum256(b); return hex.EncodeToString(s[:]) }

func TestComfyExtraArgsOnlyTakesTheAllowedSwitches(t *testing.T) {
	h, e := hostTestRunner(t)
	p := h.paths("comfyui")
	if err := os.MkdirAll(p.dir, 0o755); err != nil {
		t.Fatal(err)
	}
	file := filepath.Join(p.dir, "engine-args.txt")
	if err := os.WriteFile(file, []byte("--disable-dynamic-vram\n--listen 0.0.0.0 --disable-fast-disk --extra-model-paths-config /etc/passwd\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	argv, _ := engineArgv(h.cat, e, p)
	line := " " + strings.Join(argv, " ") + " "
	if !strings.Contains(line, " --disable-dynamic-vram ") || !strings.Contains(line, " --disable-fast-disk ") {
		t.Errorf("the allowed switches are missing: %s", line)
	}
	if strings.Contains(line, "0.0.0.0") || strings.Contains(line, "passwd") || strings.Contains(line, "--extra-model-paths-config") {
		t.Errorf("a switch outside the list got through: %s", line)
	}
	if strings.Count(line, "--listen") != 1 {
		t.Errorf("the listen address must stay as it is: %s", line)
	}
	os.Remove(file)
	argv, _ = engineArgv(h.cat, e, p)
	if strings.Contains(strings.Join(argv, " "), "--disable-dynamic-vram") {
		t.Errorf("without the file there are no extra switches")
	}
}