package provision

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

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

var ubuntu = discovery.Host{OSID: "ubuntu", OSVersion: "24.04", Arch: "amd64", CPUFlags: []string{"avx2"}}

// release is a fake GitHub: one release, its assets and their digests.
type release struct {
	assets  map[string][]byte
	digests map[string]string
	gets    int
	mu      sync.Mutex
}

func tarGz(t *testing.T, files map[string]string, links map[string]string) []byte {
	t.Helper()
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	for name, body := range files {
		_ = tw.WriteHeader(&tar.Header{Name: name, Mode: 0o755, Size: int64(len(body)), Typeflag: tar.TypeReg})
		_, _ = tw.Write([]byte(body))
	}
	for name, target := range links {
		_ = tw.WriteHeader(&tar.Header{Name: name, Linkname: target, Typeflag: tar.TypeSymlink})
	}
	_ = tw.Close()
	_ = gz.Close()
	return buf.Bytes()
}

func (r *release) serve(t *testing.T) *httptest.Server {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		r.mu.Lock()
		r.gets++
		r.mu.Unlock()
		if strings.HasPrefix(req.URL.Path, "/repos/") {
			var out struct {
				Assets []map[string]string `json:"assets"`
			}
			for name := range r.assets {
				out.Assets = append(out.Assets, map[string]string{"name": name, "digest": r.digests[name]})
			}
			_ = json.NewEncoder(w).Encode(out)
			return
		}
		body, ok := r.assets[filepath.Base(req.URL.Path)]
		if !ok {
			http.NotFound(w, req)
			return
		}
		_, _ = w.Write(body)
	}))
	t.Cleanup(srv.Close)
	return srv
}

func newRelease(assets map[string][]byte) *release {
	r := &release{assets: assets, digests: map[string]string{}}
	for name, b := range assets {
		sum := sha256.Sum256(b)
		r.digests[name] = "sha256:" + hex.EncodeToString(sum[:])
	}
	return r
}

// fakeSystem answers the package tools, ldd and the engine itself.
type fakeSystem struct {
	mu        sync.Mutex
	installed map[string]bool
	tools     map[string]bool
	lddOut    string
	// devices is what llama-server --list-devices prints, and backendLdd
	// what ldd says of its GPU backend.
	devices    string
	backendLdd string
	calls      []string
	// locked is how many more package commands find the lock held.
	locked int
}

var errExit = errors.New("exit 1")

func (f *fakeSystem) run(_ context.Context, _ []string, name string, args ...string) ([]byte, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls = append(f.calls, name+" "+strings.Join(args, " "))
	switch {
	case name == "sh" && strings.Contains(args[1], "dpkg-query"),
		name == "rpm" && args[0] == "-q":
		if f.installed[args[len(args)-1]] {
			return nil, nil
		}
		return nil, errExit
	case name == "sh" && strings.Contains(args[1], "command -v"):
		if f.tools[args[2]] {
			return nil, nil
		}
		return nil, errExit
	case name == "apt-get" || name == "dnf":
		if f.locked > 0 {
			f.locked--
			return []byte("E: Unable to lock directory /var/lib/apt/lists/\n"), errExit
		}
		for i, a := range args {
			if a == "install" {
				for _, p := range args[i+1:] {
					if !strings.HasPrefix(p, "-") {
						f.installed[p] = true
					}
				}
			}
		}
		return nil, nil
	case name == "ldd" && strings.HasSuffix(args[0], "libggml-cuda.so"):
		return []byte(f.backendLdd), nil
	case name == "ldd":
		return []byte(f.lddOut), nil
	case strings.HasSuffix(name, "llama-server") && args[0] == "--list-devices":
		return []byte(f.devices), nil
	case strings.HasSuffix(name, "llama-server"):
		return []byte("version: 10991 (abc)\n"), nil
	}
	return nil, errors.New("unexpected command " + name)
}

func installer(t *testing.T, srv *httptest.Server, sys *fakeSystem, host discovery.Host) (*Installer, *[]Status) {
	var seen []Status
	var mu sync.Mutex
	return &Installer{
		Root: t.TempDir(), Host: host, Run: sys.run,
		APIBase: srv.URL, DownloadBase: srv.URL, HTTP: srv.Client(),
		OnStatus: func(s Status) { mu.Lock(); seen = append(seen, s); mu.Unlock() },
	}, &seen
}

func llamaRelease(t *testing.T) *release {
	return newRelease(map[string][]byte{
		"llama-b1-bin-ubuntu-x64.tar.gz": tarGz(t,
			map[string]string{"llama-b1/llama-server": "#!/bin/sh\n", "llama-b1/libllama.so.0": "x"},
			map[string]string{"llama-b1/libllama.so": "libllama.so.0"}),
	})
}

func TestAnEngineIsInstalledFromNothing(t *testing.T) {
	// A node rebuilt from a bare image: no engine, no OpenMP runtime. What it
	// ends with has to have come from the publisher, checked, and been run.
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}}
	in, seen := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if len(got) != 1 || got[0].State != Ready {
		t.Fatalf("not ready: %+v", got)
	}
	if !sys.installed["libgomp1"] {
		t.Error("the OpenMP runtime the build links against was not installed")
	}
	if !strings.HasPrefix(got[0].Binary, in.Root) {
		t.Errorf("binary %s is not the one installed under %s", got[0].Binary, in.Root)
	}
	if _, err := os.Lstat(filepath.Join(filepath.Dir(got[0].Binary), "libllama.so")); err != nil {
		t.Errorf("the library symlink was not unpacked: %v", err)
	}
	var states []string
	for _, s := range *seen {
		states = append(states, s.State+": "+s.Detail)
	}
	if !strings.Contains(strings.Join(states, "\n"), "downloading") {
		t.Errorf("a download was not reported while it ran:\n%s", strings.Join(states, "\n"))
	}
}

func TestAnInstalledEngineIsCheckedNotFetchedAgain(t *testing.T) {
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)
	w := []Want{{Engine: "llamacpp", Version: "b1"}}
	in.Converge(t.Context(), w)
	before := rel.gets

	got := in.Converge(t.Context(), w)
	if got[0].State != Ready {
		t.Fatalf("second pass: %+v", got[0])
	}
	if rel.gets != before {
		t.Errorf("a ready engine was downloaded again (%d requests)", rel.gets-before)
	}
}

func TestAnEngineThatCannotLoadALibraryIsNotReady(t *testing.T) {
	// Found on a node before: llama-server present, executable, and unable to
	// start for want of a library — while the node reported ready.
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true},
		lddOut: "\tlibcuda.so.1 => not found\n"}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "libcuda.so.1") {
		t.Fatalf("expected a failure naming the library: %+v", got[0])
	}
	if _, err := os.Stat(filepath.Join(in.Root, "llamacpp", "b1-cpu", marker)); err == nil {
		t.Error("a build that does not run was marked installed")
	}
}

func TestADownloadThatDoesNotMatchItsChecksumIsRefused(t *testing.T) {
	rel := llamaRelease(t)
	for name := range rel.digests {
		rel.digests[name] = "sha256:" + strings.Repeat("0", 64)
	}
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true}}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "checksum") {
		t.Fatalf("a tampered download was not refused: %+v", got[0])
	}
}

func TestAnArchiveCannotWriteOutsideItsDirectory(t *testing.T) {
	rel := newRelease(map[string][]byte{
		"llama-b1-bin-ubuntu-x64.tar.gz": tarGz(t,
			map[string]string{"../../escaped": "x", "llama-server": ""}, nil),
	})
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true}}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed {
		t.Fatalf("an archive escaping its directory was installed: %+v", got[0])
	}
}

func TestTheBuildFollowsTheHardware(t *testing.T) {
	gpu := ubuntu
	gpu.NVIDIADriver = "570.86.15"
	got, err := llamaCppCandidates("b1", "", gpu)
	if err != nil || got[0].Variant != "cuda-12.8" || len(got[0].Assets) != 2 || got[len(got)-1].Variant != "vulkan" {
		t.Errorf("an NVIDIA node on driver 570 got %+v, %v; expected CUDA 12.8 with its runtime, then Vulkan", got, err)
	}
	if _, err := llamaCppCandidates("b1", "", discovery.Host{OSID: "ubuntu", Arch: "riscv64"}); err == nil {
		t.Error("a build was chosen for an architecture llama.cpp does not publish for")
	}
	for _, c := range []struct {
		name string
		host discovery.Host
		want []string
	}{
		// Driver 580 and Turing: CUDA 13.3 first.
		{"current", discovery.Host{Arch: "amd64", NVIDIADriver: "580.65.06", MinComputeCapability: 8.0}, []string{"cuda-13.3", "cuda-12.8", "vulkan"}},
		// A V100 on driver 550: CUDA 12.8 ships 7.0 as PTX, which needs 570.
		{"V100, 550", discovery.Host{Arch: "amd64", NVIDIADriver: "550.54.15", MinComputeCapability: 7.0}, []string{"vulkan"}},
		// An RTX 3090 on 550: 8.6 is machine code in the CUDA 12.8 build.
		{"3090, 550", discovery.Host{Arch: "amd64", NVIDIADriver: "550.54.15", MinComputeCapability: 8.6}, []string{"cuda-12.8", "vulkan"}},
		// A Pascal card: CUDA 12.8, not 13.3.
		{"P40", discovery.Host{Arch: "amd64", NVIDIADriver: "575.51.03", MinComputeCapability: 6.1}, []string{"cuda-12.8", "vulkan"}},
		// arm64 with a GPU: never the CPU build.
		{"GH200", discovery.Host{Arch: "arm64", NVIDIADriver: "580.65.06", MinComputeCapability: 9.0}, []string{"cuda-13.3", "vulkan"}},
	} {
		got, err := llamaCppCandidates("b1", "", c.host)
		var names []string
		for _, b := range got {
			names = append(names, b.Variant)
		}
		if err != nil || strings.Join(names, " ") != strings.Join(c.want, " ") {
			t.Errorf("%s: %v, %v; want %v", c.name, names, err, c.want)
		}
	}

	slow := ubuntu
	slow.CPUFlags = nil
	st := (&Installer{Host: slow}).ready(Status{}, Build{CPUFlags: []string{"avx2"}}, "/x/llama-server", "")
	if !strings.Contains(st.Detail, "avx2") {
		t.Errorf("a CPU without AVX2 was not called out: %q", st.Detail)
	}
}

func TestAnEngineTheAgentCannotInstallSaysSo(t *testing.T) {
	in := &Installer{Root: t.TempDir(), Host: ubuntu}
	got := in.Converge(t.Context(), []Want{{Engine: "tgi", Version: "3.0.0"}})
	if got[0].State != Unsupported || !strings.Contains(got[0].Detail, "install it by hand") {
		t.Errorf("got %+v", got[0])
	}
}

func TestARedHatFamilyNodeInstallsWithItsOwnPackageManager(t *testing.T) {
	// The build is not Ubuntu's in any way that matters; the package that
	// carries OpenMP's runtime is named differently, and installed by dnf.
	rocky := discovery.Host{OSID: "rocky", OSVersion: "10.0", OSLike: []string{"rhel", "centos", "fedora"},
		Arch: "amd64", CPUFlags: []string{"avx2"}}
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}, tools: map[string]bool{"dnf": true}}
	in, _ := installer(t, rel.serve(t), sys, rocky)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Ready {
		t.Fatalf("not ready: %+v", got[0])
	}
	if !sys.installed["libgomp"] || sys.installed["libgomp1"] {
		t.Errorf("installed %v; expected the rpm name, libgomp", sys.installed)
	}
}

func TestASystemTooOldForTheBuildSaysWhatItLacks(t *testing.T) {
	// RHEL 9: glibc is new enough, libstdc++ is one GCC release short. ldd
	// finds the library and reports the symbol, and "libstdc++ not found"
	// would send somebody looking for a package that is installed.
	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true},
		lddOut: "\t/lib64/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /x/libllama.so)\n" +
			"\t/lib64/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /x/llama-server)\n"}
	in, _ := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "libstdc++.so.6 is too old (needs GLIBCXX_3.4.30)") {
		t.Fatalf("detail %q", got[0].Detail)
	}
	if strings.Count(got[0].Detail, "GLIBCXX") != 1 {
		t.Errorf("the same shortfall was listed more than once: %q", got[0].Detail)
	}
}

func TestAPackageManagerBusyWithItsOwnUpdatesIsWaitedFor(t *testing.T) {
	// A new node's first minutes: cloud-init and unattended-upgrades hold
	// apt's lists lock, which DPkg::Lock::Timeout does not cover.
	lockRetry = time.Millisecond
	t.Cleanup(func() { lockRetry = 10 * time.Second })

	rel := llamaRelease(t)
	sys := &fakeSystem{installed: map[string]bool{}, locked: 3}
	in, seen := installer(t, rel.serve(t), sys, ubuntu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Ready {
		t.Fatalf("gave up on a lock another apt held: %+v", got[0])
	}
	waited := false
	for _, s := range *seen {
		waited = waited || strings.Contains(s.Detail, "waiting for another apt-get")
	}
	if !waited {
		t.Error("the wait was not reported")
	}
}

func TestAVersionThatIsAPathIsRefused(t *testing.T) {
	// The version names a directory under the install root; "../" would put
	// an install, and the RemoveAll that replaces one, outside it, as root.
	in := &Installer{Root: t.TempDir(), Host: ubuntu}
	for _, v := range []string{"../../etc", "..", ".hidden", "b1/../../x", "", "b1 && rm -rf /"} {
		got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: v}})
		if got[0].State != Unsupported {
			t.Errorf("version %q: %+v", v, got[0])
		}
	}
}

func TestChainedLinksCannotCarryAFileOutOfTheInstall(t *testing.T) {
	// Each link reads as inside the directory; together they point above it.
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	_ = tw.WriteHeader(&tar.Header{Name: "a", Linkname: ".", Typeflag: tar.TypeSymlink})
	_ = tw.WriteHeader(&tar.Header{Name: "a/a/b", Linkname: "../..", Typeflag: tar.TypeSymlink})
	_ = tw.WriteHeader(&tar.Header{Name: "a/a/b/escaped", Mode: 0o644, Size: 1, Typeflag: tar.TypeReg})
	_, _ = tw.Write([]byte("x"))
	_ = tw.Close()
	_ = gz.Close()

	parent := t.TempDir()
	dir := filepath.Join(parent, "install")
	_ = os.MkdirAll(dir, 0o755)
	archive := filepath.Join(parent, "evil.tar.gz")
	_ = os.WriteFile(archive, buf.Bytes(), 0o644)

	if err := untarGz(archive, dir); err == nil {
		t.Error("an archive of chained links was unpacked")
	}
	if _, err := os.Stat(filepath.Join(parent, "escaped")); err == nil {
		t.Fatal("a file was written outside the install directory")
	}
}

func TestALibrarysVersionLinkIsStillUnpacked(t *testing.T) {
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	_ = tw.WriteHeader(&tar.Header{Name: "lib/libllama.so.0", Mode: 0o755, Size: 1, Typeflag: tar.TypeReg})
	_, _ = tw.Write([]byte("x"))
	_ = tw.WriteHeader(&tar.Header{Name: "lib/libllama.so", Linkname: "libllama.so.0", Typeflag: tar.TypeSymlink})
	_ = tw.Close()
	_ = gz.Close()
	dir := t.TempDir()
	archive := filepath.Join(t.TempDir(), "ok.tar.gz")
	_ = os.WriteFile(archive, buf.Bytes(), 0o644)
	if err := untarGz(archive, dir); err != nil {
		t.Fatal(err)
	}
	if target, err := os.Readlink(filepath.Join(dir, "lib", "libllama.so")); err != nil || target != "libllama.so.0" {
		t.Errorf("link %q %v", target, err)
	}
}

func TestACUDABuildThatSeesNoGPUIsNotInstalled(t *testing.T) {
	// Found in the field: libggml-cuda.so in the build, cuBLAS nowhere the
	// loader looks, and llama-server serving a GPU placement on the CPU.
	cuda := "bin-ubuntu-cuda-12.8-x64.tar.gz"
	rel := newRelease(map[string][]byte{
		"llama-b1-" + cuda: tarGz(t, map[string]string{
			"llama-b1/llama-server": "#!/bin/sh\n", "llama-b1/libggml-cuda.so": "x"}, nil),
		"cudart-llama-b1-" + cuda: tarGz(t, map[string]string{"libcudart.so.12": "x"}, nil),
	})
	gpu := ubuntu
	gpu.NVIDIADriver = "570.86.15"
	sys := &fakeSystem{installed: map[string]bool{"libgomp1": true},
		devices:    "Available devices:\n  (none)\n",
		backendLdd: "\tlibcublas.so.12 => not found\n\tlibcudart.so.12 => /x/libcudart.so.12\n"}
	in, _ := installer(t, rel.serve(t), sys, gpu)

	got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}})
	if got[0].State != Failed || !strings.Contains(got[0].Detail, "sees no GPU") ||
		!strings.Contains(got[0].Detail, "libcublas.so.12 is not installed") {
		t.Fatalf("expected a failure naming the missing library: %+v", got[0])
	}

	// The same build seeing its GPU is installed.
	sys.devices = "Available devices:\n  CUDA0: NVIDIA L4 (22994 MiB, 22713 MiB free)\n"
	if got := in.Converge(t.Context(), []Want{{Engine: "llamacpp", Version: "b1"}}); got[0].State != Ready {
		t.Fatalf("a CUDA build that sees its GPU: %+v", got[0])
	}
}