package main

import (
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func fixture(t *testing.T) (Config, string) {
	root := t.TempDir()
	dir := filepath.Join(root, "nvidia", "Qwen3.8-27B-NVFP4")
	os.MkdirAll(dir, 0755)
	os.WriteFile(filepath.Join(dir, "config.json"), []byte("{}"), 0644)
	cfg := defaultConfig()
	cfg.ModelsDir = root
	cfg.AgentUID = 1000
	return cfg, root
}

func goodReq() Request {
	return Request{Op: "start", Model: "nvidia/Qwen3.8-27B-NVFP4", Port: 8100, Mem: 30 << 30, Util: 0.3, MaxLen: 32768}
}

func TestBuildRunProducesTheExpectedLockedDownCommand(t *testing.T) {
	cfg, root := fixture(t)
	args, err := buildRun(cfg, goodReq(), nil)
	if err != nil {
		t.Fatal(err)
	}
	joined := strings.Join(args, " ")
	for _, want := range []string{"--gpus all", "--security-opt no-new-privileges", "--cap-drop ALL", "127.0.0.1:8100:8000", cfg.Image, "vllm serve /models/nvidia/Qwen3.8-27B-NVFP4", "--gpu-memory-utilization 0.30", "HF_HUB_OFFLINE=1"} {
		if !strings.Contains(joined, want) {
			t.Errorf("missing %q in: %s", want, joined)
		}
	}
	for _, bad := range []string{"--privileged", "--network host", "--net=host", "--ipc=host", "--pid"} {
		if strings.Contains(joined, bad) {
			t.Errorf("must not contain %q", bad)
		}
	}
	rootReal, _ := filepath.EvalSymlinks(root)
	if !strings.Contains(joined, rootReal+":/models:ro") {
		t.Errorf("models must be mounted read-only: %s", joined)
	}
}

func TestBuildRunRejectsBadInputs(t *testing.T) {
	cfg, _ := fixture(t)
	cases := map[string]func(*Request){
		"path traversal in name":  func(r *Request) { r.Model = "../../etc/passwd" },
		"name without org":        func(r *Request) { r.Model = "justaname" },
		"shell characters":        func(r *Request) { r.Model = "a/b;rm -rf" },
		"model that is not there": func(r *Request) { r.Model = "nvidia/Other" },
		"port too low":            func(r *Request) { r.Port = 22 },
		"port too high":           func(r *Request) { r.Port = 9000 },
		"utilisation too high":    func(r *Request) { r.Util = 0.99 },
		"utilisation too low":     func(r *Request) { r.Util = 0 },
		"tiny context":            func(r *Request) { r.MaxLen = 10 },
		"unknown flag":            func(r *Request) { r.Args = []string{"--privileged"} },
		"volume smuggling":        func(r *Request) { r.Args = []string{"-v", "/:/host"} },
		"bad flag value":          func(r *Request) { r.Args = []string{"--kv-cache-dtype", "rm -rf"} },
		"flag without value":      func(r *Request) { r.Args = []string{"--kv-cache-dtype"} },
		"trust remote code":       func(r *Request) { r.Args = []string{"--trust-remote-code"} },
	}
	for name, mutate := range cases {
		r := goodReq()
		mutate(&r)
		if _, err := buildRun(cfg, r, nil); err == nil {
			t.Errorf("%s: should have been rejected", name)
		}
	}
}

func TestBuildRunAcceptsAllowListedFlags(t *testing.T) {
	cfg, _ := fixture(t)
	r := goodReq()
	r.Args = []string{"--kv-cache-dtype", "fp8", "--enable-prefix-caching", "--max-num-seqs", "8", "--reasoning-parser", "qwen3"}
	args, err := buildRun(cfg, r, nil)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.HasSuffix(strings.Join(args, " "), "--kv-cache-dtype fp8 --enable-prefix-caching --max-num-seqs 8 --reasoning-parser qwen3") {
		t.Fatalf("%v", args)
	}
}

func TestPortAndContainerLimits(t *testing.T) {
	cfg, _ := fixture(t)
	existing := []container{{Name: "wxch-vllm-a--b", Model: "a/b", Port: 8100, State: "running"}}
	if _, err := buildRun(cfg, goodReq(), existing); err == nil {
		t.Fatal("a port that is already used must be refused")
	}
	r := goodReq()
	r.Port = 8101
	existing = []container{{Name: "wxch-vllm-x--1", Port: 8100}, {Name: "wxch-vllm-x--2", Port: 8102}}
	if _, err := buildRun(cfg, r, existing); err == nil {
		t.Fatal("container limit must be enforced")
	}
	// restarting the same model is allowed even at the limit
	existing = []container{{Name: containerOf(r.Model), Port: 8101}, {Name: "wxch-vllm-x--2", Port: 8102}}
	if _, err := buildRun(cfg, r, existing); err != nil {
		t.Fatalf("re-starting the same model must be allowed: %v", err)
	}
}

func TestSymlinkCannotEscapeTheModelsFolder(t *testing.T) {
	cfg, root := fixture(t)
	outside := t.TempDir()
	os.WriteFile(filepath.Join(outside, "config.json"), []byte("{}"), 0644)
	link := filepath.Join(root, "evil", "model")
	os.MkdirAll(filepath.Dir(link), 0755)
	if err := os.Symlink(outside, link); err != nil {
		t.Skip("cannot create symlinks here: " + err.Error())
	}
	r := goodReq()
	r.Model = "evil/model"
	if _, err := buildRun(cfg, r, nil); err == nil {
		t.Fatal("a symlink pointing outside the models folder must be refused")
	}
}

func TestUnknownOperationAndNames(t *testing.T) {
	cfg, _ := fixture(t)
	if r := handle(cfg, Request{Op: "exec", Model: "a/b"}); r.OK {
		t.Fatal("unknown operations must fail")
	}
	if r := handle(cfg, Request{Op: "stop", Model: "x; reboot"}); r.OK {
		t.Fatal("bad names must fail")
	}
	if r := handle(cfg, Request{Op: "logs", Model: "../x"}); r.OK {
		t.Fatal("bad names must fail")
	}
}

func TestModelLinkedIntoAListedExtraRootIsMountedAtTheSamePath(t *testing.T) {
	cfg, root := fixture(t)
	disk := t.TempDir()
	real := filepath.Join(disk, "org", "model")
	os.MkdirAll(real, 0755)
	os.WriteFile(filepath.Join(real, "config.json"), []byte("{}"), 0644)
	link := filepath.Join(root, "org", "model")
	os.MkdirAll(filepath.Dir(link), 0755)
	if err := os.Symlink(real, link); err != nil {
		t.Skip("cannot create symlinks here: " + err.Error())
	}
	r := goodReq()
	r.Model = "org/model"
	if _, err := buildRun(cfg, r, nil); err == nil {
		t.Fatal("a link out of the models folder must be refused unless root listed the folder")
	}
	cfg.ExtraRoots = []string{disk}
	args, err := buildRun(cfg, r, nil)
	if err != nil {
		t.Fatal(err)
	}
	diskReal, _ := filepath.EvalSymlinks(disk)
	joined := strings.Join(args, " ")
	if !strings.Contains(joined, "-v "+diskReal+":"+diskReal+":ro") {
		t.Fatalf("the disk folder must be mounted read-only at its own path: %s", joined)
	}
	if strings.Index(joined, diskReal+":"+diskReal) > strings.Index(joined, cfg.Image) {
		t.Fatalf("mounts must come before the image name: %s", joined)
	}
	// a listed root does not open anything else
	other := t.TempDir()
	os.WriteFile(filepath.Join(other, "config.json"), []byte("{}"), 0644)
	evil := filepath.Join(root, "evil", "model")
	os.MkdirAll(filepath.Dir(evil), 0755)
	os.Symlink(other, evil)
	r.Model = "evil/model"
	if _, err := buildRun(cfg, r, nil); err == nil {
		t.Fatal("only the listed extra root may be linked to")
	}
}

func TestConfigIsReadAgainSoANewDiskNeedsNoRestart(t *testing.T) {
	path := filepath.Join(t.TempDir(), "runner.json")
	os.WriteFile(path, []byte(`{"agent_uid": 1000, "models_dir": "/m"}`), 0600)
	c, err := readConfig(path)
	if err != nil || len(c.ExtraRoots) != 0 || c.AgentUID != 1000 {
		t.Fatalf("%+v %v", c, err)
	}
	os.WriteFile(path, []byte(`{"agent_uid": 1000, "models_dir": "/m", "extra_roots": ["/mnt/wxch-ext"]}`), 0600)
	if c, err = readConfig(path); err != nil || len(c.ExtraRoots) != 1 || c.ExtraRoots[0] != "/mnt/wxch-ext" {
		t.Fatalf("%+v %v", c, err)
	}
	os.WriteFile(path, []byte(`{"models_dir": "/m"}`), 0600)
	if _, err = readConfig(path); err == nil {
		t.Fatal("a config without agent_uid must be refused")
	}
	if _, err = readConfig(filepath.Join(t.TempDir(), "missing.json")); err == nil {
		t.Fatal("a missing config is an error")
	}
}

func TestOnlyTheAgentProgramMayUseTheRunnerWhenItIsConfigured(t *testing.T) {
	cfg := Config{AgentUID: 1000, AgentExe: "/opt/wxch-agent/wxch-agent"}
	exe := map[int]string{10: "/opt/wxch-agent/wxch-agent", 11: "/home/pgx/comfy/venv/bin/python3.12", 12: "/usr/bin/bash"}
	exeOf := func(pid int) (string, error) {
		if p, ok := exe[pid]; ok {
			return p, nil
		}
		return "", os.ErrNotExist
	}
	cases := []struct {
		name     string
		uid, pid int
		want     bool
	}{
		{"the agent", 1000, 10, true},
		{"an engine running as the same user", 1000, 11, false},
		{"a shell of the same user", 1000, 12, false},
		{"a process that is gone", 1000, 99, false},
		{"another user", 1001, 10, false},
		{"root (the installer)", 0, 12, true},
	}
	for _, c := range cases {
		if got := peerAllowed(cfg, c.uid, c.pid, exeOf); got != c.want {
			t.Errorf("%s: got %v, want %v", c.name, got, c.want)
		}
	}
	// an older config without agent_exe keeps checking the user only
	old := Config{AgentUID: 1000}
	if !peerAllowed(old, 1000, 11, exeOf) || peerAllowed(old, 1001, 10, exeOf) {
		t.Error("without agent_exe only the user is checked")
	}
}