package provision

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
)

var gpuNode = discovery.Host{OSID: "ubuntu", OSVersion: "24.04", Arch: "amd64",
	NVIDIADriver: "580.65.06", MinComputeCapability: 8.0}

// fakePython is uv and the environments it builds.
type fakePython struct {
	fakeSystem
	pmu      sync.Mutex
	commands []string
	failPip  bool
	// noKernel makes the GPU check fail as a torch without kernels for the
	// card does.
	noKernel bool
}

func (f *fakePython) run(ctx context.Context, env []string, name string, args ...string) ([]byte, error) {
	line := name + " " + strings.Join(args, " ")
	f.pmu.Lock()
	f.commands = append(f.commands, line)
	f.pmu.Unlock()
	switch {
	case strings.HasSuffix(name, "/uv") && len(args) == 1 && args[0] == "--version":
		return []byte("uv 0.12.20\n"), nil
	case strings.HasSuffix(name, "/uv") && args[0] == "venv":
		return nil, os.MkdirAll(filepath.Join(args[len(args)-1], "bin"), 0o755)
	case strings.HasSuffix(name, "/uv") && args[0] == "pip" && args[1] == "install":
		if f.failPip {
			return []byte("error: Hash mismatch for torch\n"), errExit
		}
		venv := filepath.Dir(filepath.Dir(args[slices.Index(args, "--python")+1]))
		return nil, os.WriteFile(filepath.Join(venv, "bin", "vllm"), []byte("#!/bin/sh\n"), 0o755)
	case strings.HasSuffix(name, "/bin/vllm") && args[0] == "--version":
		return []byte("0.30.0\n"), nil
	case strings.HasSuffix(name, "/bin/python") && args[0] == "-c":
		if f.noKernel {
			return []byte("RuntimeError: CUDA error: no kernel image is available for execution on the device\n"), errExit
		}
		return []byte("128.0 Tesla V100-SXM2-32GB (7, 0)\n"), nil
	}
	return f.fakeSystem.run(ctx, env, name, args...)
}

func (f *fakePython) runStdin(_ context.Context, _ []string, stdin, name string, args ...string) ([]byte, error) {
	f.pmu.Lock()
	f.commands = append(f.commands, name+" "+strings.Join(args, " ")+" <<< "+strings.TrimSpace(stdin))
	f.pmu.Unlock()
	out := args[slices.Index(args, "--output-file")+1]
	return nil, os.WriteFile(out, []byte("vllm==0.30.0 --hash=sha256:ab\n"), 0o644)
}

func pythonInstaller(t *testing.T, host discovery.Host, sys *fakePython) (*Installer, *release) {
	uv := newRelease(map[string][]byte{
		"uv-x86_64-unknown-linux-gnu.tar.gz": tarGz(t,
			map[string]string{"uv-x86_64-unknown-linux-gnu/uv": "#!/bin/sh\n"}, nil),
		// The fake GitHub serves every repository's release from this one.
		"vllm-0.30.0+cu129-cp38-abi3-manylinux_2_28_x86_64.whl": []byte("the cu129 wheel"),
	})
	gh := uv.serve(t)
	mux := http.NewServeMux()
	mux.HandleFunc("/pypi/vllm/0.30.0/json", func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"urls":[{"upload_time_iso_8601":"2026-09-22T05:22:21.647189Z"},
			{"upload_time_iso_8601":"2026-09-22T05:23:10.039243Z"}]}`))
	})
	pypi := httptest.NewServer(mux)
	t.Cleanup(pypi.Close)
	return &Installer{
		Root: t.TempDir(), Host: host, Run: sys.run, RunStdin: sys.runStdin,
		APIBase: gh.URL, DownloadBase: gh.URL, PyPIBase: pypi.URL, HTTP: gh.Client(),
		FreeSpace: func(string) (float64, error) { return 100, nil },
	}, uv
}

func TestVLLMIsInstalledOnAGPUNode(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, _ := pythonInstaller(t, gpuNode, sys)

	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Ready {
		t.Fatalf("not ready: %+v", got[0])
	}
	dir := filepath.Join(in.Root, "vllm", "0.30.0-cu13")
	if got[0].Binary != filepath.Join(dir, "bin", "vllm") {
		t.Errorf("binary %s, want the environment's vllm", got[0].Binary)
	}
	if !sys.installed["gcc"] {
		t.Error("gcc, which Triton compiles its launcher with, was not installed")
	}
	all := strings.Join(sys.commands, "\n")
	for _, want := range []string{
		// PyPI as it stood a day after vLLM's last file, so the dependencies
		// are the ones it was released against.
		"--exclude-newer 2026-09-23T05:23:10Z",
		"--generate-hashes",
		"--require-hashes",
		"venv --relocatable --python 3.12",
	} {
		if !strings.Contains(all, want) {
			t.Errorf("no %q in:\n%s", want, all)
		}
	}
	if _, err := os.Stat(filepath.Join(dir, lockName)); err != nil {
		t.Errorf("the resolved requirements were not kept: %v", err)
	}
	if _, err := os.Stat(filepath.Join(in.Root, ".uv-cache")); err == nil {
		t.Error("the download cache was left behind after a successful install")
	}

	// Checked again on the next pass, not installed again.
	before := len(sys.commands)
	if again := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}}); again[0].State != Ready {
		t.Fatalf("second pass: %+v", again[0])
	}
	for _, c := range sys.commands[before:] {
		if strings.Contains(c, "pip install") {
			t.Errorf("a ready environment was installed again: %s", c)
		}
	}
}

func TestVLLMIsNotForEveryNode(t *testing.T) {
	cases := []struct {
		name  string
		host  discovery.Host
		state string
		says  string
	}{
		{"no GPU", ubuntu, Skipped, "has none"},
		// Too old for CUDA 13 and for its CUDA 12.9 build both.
		{"a driver too old for either", discovery.Host{Arch: "amd64", NVIDIADriver: "470.256.02"}, Unsupported, "driver 525"},
		{"Pascal", discovery.Host{Arch: "amd64", NVIDIADriver: "580.65.06", MinComputeCapability: 6.1}, Unsupported, "7.0 or newer"},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
			in, uv := pythonInstaller(t, c.host, sys)
			got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
			if got[0].State != c.state || !strings.Contains(got[0].Detail, c.says) {
				t.Errorf("got %s %q, want %s saying %q", got[0].State, got[0].Detail, c.state, c.says)
			}
			if uv.gets != 0 || len(sys.commands) != 0 {
				t.Errorf("a node that cannot run it still started installing: %v", sys.commands)
			}
		})
	}
}

func TestAPythonEngineIsRefusedWithoutTheDiskForIt(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{"gcc": true}}}
	in, uv := pythonInstaller(t, gpuNode, sys)
	in.FreeSpace = func(string) (float64, error) { return 12.5, nil }
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "12.5 GB") {
		t.Errorf("got %s %q", got[0].State, got[0].Detail)
	}
	if uv.gets != 0 {
		t.Error("it began downloading anyway")
	}
}

func TestAFailedPythonInstallLeavesNothingThatLooksInstalled(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{"gcc": true}}, failPip: true}
	in, _ := pythonInstaller(t, gpuNode, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "Hash mismatch") {
		t.Fatalf("got %s %q", got[0].State, got[0].Detail)
	}
	entries, _ := os.ReadDir(filepath.Join(in.Root, "vllm"))
	for _, e := range entries {
		t.Errorf("left behind: %s", e.Name())
	}
}

func TestAVersionPyPIDoesNotHaveIsSaidSo(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{"gcc": true}}}
	in, _ := pythonInstaller(t, gpuNode, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "9.9.9"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "PyPI has no vllm 9.9.9") {
		t.Errorf("got %s %q", got[0].State, got[0].Detail)
	}
}

func TestAnOlderGPUGetsTheSameVLLMBuiltForCUDA12(t *testing.T) {
	// A V100: CUDA 13 dropped it. The same vLLM's CUDA 12.9 build, from the
	// project's GitHub release, with the torch PyTorch publishes for 12.9.
	volta := discovery.Host{OSID: "ubuntu", Arch: "amd64", NVIDIADriver: "550.54.15", MinComputeCapability: 7.0}
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, gh := pythonInstaller(t, volta, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Ready || got[0].Variant != "cu129" {
		t.Fatalf("a V100 node: %+v", got[0])
	}
	if _, err := os.Stat(filepath.Join(in.Root, "vllm", "0.30.0-cu129", marker)); err != nil {
		t.Errorf("not installed under its own build's name: %v", err)
	}
	all := strings.Join(sys.commands, "\n")
	wheel := "vllm-0.30.0+cu129-cp38-abi3-manylinux_2_28_x86_64.whl"
	for _, want := range []string{
		// Pinned to the digest GitHub publishes for the wheel.
		"vllm @ " + in.DownloadBase + "/vllm-project/vllm/releases/download/v0.30.0/" + wheel + "#sha256=" +
			strings.TrimPrefix(gh.digests[wheel], "sha256:"),
		"--extra-index-url https://download.pytorch.org/whl/cu129 --index-strategy unsafe-best-match -",
		"--require-hashes --no-deps -r",
		// Computing on the card, not only importing.
		"/bin/python -c import torch",
	} {
		if !strings.Contains(all, want) {
			t.Errorf("no %q in:\n%s", want, all)
		}
	}
	for _, c := range sys.commands {
		if strings.Contains(c, "pip install") && !strings.Contains(c, "download.pytorch.org/whl/cu129") {
			t.Errorf("installed without the index its torch comes from: %s", c)
		}
	}
}

func TestABuildChosenForANodeIsTheOneInstalled(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}}
	in, _ := pythonInstaller(t, gpuNode, sys)
	if got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0", Variant: "cu129"}}); got[0].Variant != "cu129" || got[0].State != Ready {
		t.Errorf("cu129 chosen for a node that runs cu13 too: %+v", got[0])
	}
	volta := discovery.Host{Arch: "amd64", NVIDIADriver: "580.65.06", MinComputeCapability: 7.0}
	if _, err := chooseVLLM("0.30.0", "cu13", volta); err == nil || !strings.Contains(err.Error(), "7.5") {
		t.Errorf("cu13 chosen for a V100 should be refused, saying why: %v", err)
	}
	if _, err := chooseVLLM("0.30.0", "cu200", gpuNode); err == nil {
		t.Error("a build that does not exist was accepted")
	}
}

func TestABuildThatCannotComputeOnTheCardIsNotInstalled(t *testing.T) {
	sys := &fakePython{fakeSystem: fakeSystem{installed: map[string]bool{}}, noKernel: true}
	in, _ := pythonInstaller(t, gpuNode, sys)
	got := in.Converge(t.Context(), []Want{{Engine: "vllm", Version: "0.30.0"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "no kernels for this GPU") {
		t.Fatalf("a build without kernels for the card: %+v", got[0])
	}
	if _, err := os.Stat(filepath.Join(in.Root, "vllm", "0.30.0-cu13", marker)); err == nil {
		t.Error("marked installed")
	}
}

func TestOnlyACUDARequirementFallsBackToTheCUDA12Build(t *testing.T) {
	// A machine vLLM is not published for is refused, not given the other
	// build to fail with.
	if _, err := chooseVLLM("0.30.0", "", discovery.Host{NVIDIADriver: "580.65.06"}); err == nil ||
		!strings.Contains(err.Error(), "no Linux build") {
		t.Errorf("an unknown architecture: %v", err)
	}
	if _, err := chooseVLLM("0.30.0", "cu129", ubuntu); !errors.As(err, new(NotHere)) {
		t.Errorf("a node with no GPU chosen cu129: %v", err)
	}
}