package budget

import (
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"
)

func waitGroup(t *testing.T, p procFS, pid uint32, part string) string {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		group, err := p.Group(pid)
		if err != nil {
			t.Fatal(err)
		}
		if strings.Contains(group, part) {
			return group
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatalf("PID %d never entered %s", pid, part)
	return ""
}

func startChild(t *testing.T, command string, args ...string) *exec.Cmd {
	t.Helper()
	cmd := exec.Command(command, args...)
	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		if cmd.ProcessState != nil {
			return
		}
		if err := cmd.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) {
			t.Error(err)
		}
		var exit *exec.ExitError
		if err := cmd.Wait(); err != nil && !errors.As(err, &exit) {
			t.Error(err)
		}
	})
	return cmd
}

func TestLiveManager(t *testing.T) {
	if os.Getenv("BUDGET_LIVE") != "1" {
		t.Skip("set BUDGET_LIVE=1 for the user-manager check")
	}
	home, err := os.UserHomeDir()
	if err != nil {
		t.Fatal(err)
	}
	p := homePaths(home)
	c, err := loadConfig(p.config)
	if err != nil {
		t.Fatal(err)
	}
	files, err := c.render()
	if err != nil {
		t.Fatal(err)
	}
	for name, want := range files {
		data, err := os.ReadFile(filepath.Join(p.units, name))
		if err != nil || string(data) != want {
			t.Fatalf("generated file differs: %s, %v", name, err)
		}
	}
	m, err := connect()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		if err := m.conn.Close(); err != nil {
			t.Error(err)
		}
	})
	if _, err := m.Units(); err != nil {
		t.Fatal(err)
	}
	child := startChild(t, "sleep", "30")
	id, err := randomID()
	if err != nil {
		t.Fatal(err)
	}
	name := "app-bud-" + id + ".scope"
	if err := m.Start(name, []property{prop("PIDs", []uint32{uint32(child.Process.Pid)}), prop("Slice", "app.slice"), prop("CollectMode", "inactive-or-failed")}); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		if err := m.Stop(name); err != nil && !busError(err, noUnit) {
			t.Error(err)
		}
	})
	group := waitGroup(t, procFS{"/proc", uint32(os.Getuid())}, uint32(child.Process.Pid), "/"+name)
	root := filepath.Join("/sys/fs/cgroup", strings.TrimPrefix(group, "/"))
	props, err := c.properties("default")
	if err != nil {
		t.Fatal(err)
	}
	if err := m.Set(name, props); err != nil {
		t.Fatal(err)
	}
	got, err := os.ReadFile(filepath.Join(root, "memory.max"))
	if err != nil {
		t.Fatal(err)
	}
	want := strconv.FormatUint(values(props)["MemoryMax"].Value().(uint64), 10)
	if strings.TrimSpace(string(got)) != want {
		t.Fatalf("memory.max = %q, want %q", got, want)
	}
	state, err := m.State(name)
	if err != nil || state.Active != "active" || state.Group != group {
		t.Fatalf("%#v, %v", state, err)
	}
	props, err = c.properties("unrestricted")
	if err != nil {
		t.Fatal(err)
	}
	if err := m.Set(name, props); err != nil {
		t.Fatal(err)
	}
	for file, want := range map[string]string{"cpu.max": "max 100000", "memory.max": "max", "io.max": ""} {
		got, err := os.ReadFile(filepath.Join(root, file))
		if err != nil || strings.TrimSpace(string(got)) != want {
			t.Fatalf("%s = %q, %v; want %q", file, got, err, want)
		}
	}
	if err := m.Stop(name); err != nil {
		t.Fatal(err)
	}
}

func security(t *testing.T, pid int) map[string]string {
	t.Helper()
	data, err := os.ReadFile(filepath.Join("/proc", strconv.Itoa(pid), "status"))
	if err != nil {
		t.Fatal(err)
	}
	result := map[string]string{}
	for _, line := range strings.Split(string(data), "\n") {
		key, value, _ := strings.Cut(line, ":")
		if key == "NoNewPrivs" || key == "Seccomp" {
			result[key] = strings.TrimSpace(value)
		}
	}
	return result
}

func TestLiveWatcher(t *testing.T) {
	if os.Getenv("BUDGET_WATCH") != "1" {
		t.Skip("set BUDGET_WATCH=1 for the terminal watcher check")
	}
	root := t.TempDir()
	children := []*exec.Cmd{}
	before := []map[string]string{}
	for i := range 2 {
		fifo := filepath.Join(root, "input-"+strconv.Itoa(i))
		if err := syscall.Mkfifo(fifo, 0600); err != nil {
			t.Fatal(err)
		}
		child := startChild(t, "rg", "--no-mmap", "-q", "absent", fifo)
		children = append(children, child)
		before = append(before, security(t, child.Process.Pid))
	}
	groups := []string{}
	for _, child := range children {
		groups = append(groups, waitGroup(t, procFS{"/proc", uint32(os.Getuid())}, uint32(child.Process.Pid), "/search.slice/budget-s"))
	}
	if groups[0] == groups[1] {
		t.Fatal("searches didn't receive distinct scopes")
	}
	for i, child := range children {
		after := security(t, child.Process.Pid)
		if before[i]["NoNewPrivs"] != after["NoNewPrivs"] || before[i]["Seccomp"] != after["Seccomp"] {
			t.Fatal("sandbox security changed")
		}
		stream, err := os.OpenFile(filepath.Join(root, "input-"+strconv.Itoa(i)), os.O_WRONLY, 0600)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := stream.WriteString("present\n"); err != nil {
			t.Fatal(err)
		}
		if err := stream.Close(); err != nil {
			t.Fatal(err)
		}
		var exit *exec.ExitError
		if err := child.Wait(); !errors.As(err, &exit) || exit.ExitCode() != 1 {
			t.Fatalf("unexpected search exit: %v", err)
		}
	}
}

func TestLiveLaunch(t *testing.T) {
	if os.Getenv("BUDGET_LIVE") != "1" {
		t.Skip("set BUDGET_LIVE=1 for the native launcher check")
	}
	wd, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	cmd := exec.Command(filepath.Join("..", "..", "bin", "budget"), "run", "unrestricted", "--", "sh", "-c", `test "$(pwd -P)" = "$1" && exit 7`, "sh", wd)
	var exit *exec.ExitError
	if err := cmd.Run(); !errors.As(err, &exit) || exit.ExitCode() != 7 {
		t.Fatalf("working directory or exit status changed: %v", err)
	}
}