package tentacle

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"
)

func TestContainmentRejectsInvalidConfiguration(t *testing.T) {
	for _, change := range []func(*ResourceLimits){
		func(limits *ResourceLimits) { limits.PerSession.CPUQuota = 0 },
		func(limits *ResourceLimits) { limits.PerSession.MemoryBytes = 0 },
		func(limits *ResourceLimits) { limits.PerSession.Tasks = 0 },
		func(limits *ResourceLimits) { limits.Total.CPUQuota = limits.PerSession.CPUQuota - 1 },
		func(limits *ResourceLimits) { limits.Total.MemoryBytes = limits.PerSession.MemoryBytes - 1 },
		func(limits *ResourceLimits) { limits.Total.Tasks = limits.PerSession.Tasks - 1 },
		func(limits *ResourceLimits) { limits.CleanupTimeout = 0 },
	} {
		limits := DefaultResourceLimits()
		change(&limits)
		if _, err := NewContainedManager(DefaultLimits(), nil, t.TempDir(), limits); !errors.Is(err, ErrContainment) {
			t.Fatalf("invalid budget accepted: %v", err)
		}
	}
	// Ordinary directories must never be treated as working cgroup controllers.
	parent := t.TempDir()
	if _, err := NewContainedManager(DefaultLimits(), nil, parent, DefaultResourceLimits()); !errors.Is(err, ErrContainment) {
		t.Fatalf("ordinary directory accepted: %v", err)
	}
	entries, err := os.ReadDir(parent)
	if err != nil || len(entries) != 0 {
		t.Fatalf("failed setup mutated ordinary parent: %v %v", entries, err)
	}
}

func contained_manager(t *testing.T, resources ResourceLimits) *Manager {
	t.Helper()
	parent := os.Getenv("CTHULHU_TENTACLE_CGROUP_PARENT")
	if parent == "" {
		t.Skip("set CTHULHU_TENTACLE_CGROUP_PARENT to an empty delegated cgroup v2 parent with cpu/memory/pids enabled")
	}
	limits := DefaultLimits()
	limits.RequestTimeout = 4 * time.Second
	manager, err := NewContainedManager(limits, nil, parent, resources)
	if err != nil {
		t.Fatal(err) // An explicitly requested native check must not silently skip.
	}
	groupPath := manager.containment.group.path
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := manager.Shutdown(ctx); err != nil {
			t.Error(err)
		}
		if stats := manager.Stats(); stats != (Stats{}) {
			t.Errorf("retained resources: %+v", stats)
		}
		if _, err := os.Stat(groupPath); !os.IsNotExist(err) {
			t.Errorf("aggregate group remains: %v", err)
		}
		if err := manager.Shutdown(ctx); err != nil {
			t.Errorf("repeated shutdown: %v", err)
		}
	})
	return manager
}

func resource_peer(t *testing.T, manager *Manager, id string) *Session {
	t.Helper()
	python, err := exec.LookPath("python3")
	if err != nil {
		t.Fatal(err)
	}
	fixture, err := filepath.Abs("testdata/resource_peer.py")
	if err != nil {
		t.Fatal(err)
	}
	session, err := manager.Start(context.Background(), Spec{ID: id, Program: []string{python, "-I", "-u", fixture}})
	if err != nil {
		t.Fatal(err)
	}
	wait_signal(t, session.Ready())
	if state, err := session.Status(); state != Ready || err != nil {
		t.Fatalf("resource peer %s: %v", state, err)
	}
	if err := session.Activate(); err != nil {
		t.Fatal(err)
	}
	return session
}

func resource_call(t *testing.T, session *Session, arguments string) Result {
	t.Helper()
	pending, err := session.Invoke(context.Background(), "echo", json.RawMessage(arguments))
	if err != nil {
		t.Fatal(err)
	}
	return wait_result(t, pending)
}

func read_control(t *testing.T, path, name string) string {
	t.Helper()
	data, err := os.ReadFile(filepath.Join(path, name))
	if err != nil {
		t.Fatal(err)
	}
	return strings.TrimSpace(string(data))
}

func control_counter(t *testing.T, path, name, key string) int64 {
	t.Helper()
	for _, line := range strings.Split(read_control(t, path, name), "\n") {
		fields := strings.Fields(line)
		if len(fields) == 2 && fields[0] == key {
			value, err := strconv.ParseInt(fields[1], 10, 64)
			if err != nil {
				t.Fatal(err)
			}
			return value
		}
	}
	t.Fatalf("missing %s counter in %s", key, name)
	return 0
}

func wait_control(t *testing.T, check func() bool) {
	t.Helper()
	deadline := time.NewTimer(4 * time.Second)
	defer deadline.Stop()
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	for !check() {
		select {
		case <-ticker.C:
		case <-deadline.C:
			t.Fatal("kernel resource condition was not observed")
		}
	}
}

func TestNativeContainmentPlacementAndDetachedCleanup(t *testing.T) {
	manager := contained_manager(t, DefaultResourceLimits())
	session := resource_peer(t, manager, "detaching")
	groupPath := filepath.Join(manager.containment.group.path, session.Identity())
	for name, expected := range map[string]string{
		"cpu.max": "25000 100000", "memory.max": "268435456", "memory.swap.max": "0",
		"pids.max": "64", "memory.oom.group": "1", "cgroup.max.descendants": "0",
	} {
		if actual := read_control(t, groupPath, name); actual != expected {
			t.Fatalf("%s = %s, want %s", name, actual, expected)
		}
	}
	result := resource_call(t, session, `{}`)
	var location string
	if result.Err != nil || json.Unmarshal(result.Value, &location) != nil || !strings.HasSuffix(location, "/"+session.Identity()) {
		t.Fatalf("child ran outside its group: %+v", result)
	}
	result = resource_call(t, session, `{"action":"detach"}`)
	var pid int
	if result.Err != nil || json.Unmarshal(result.Value, &pid) != nil || pid <= 1 {
		t.Fatalf("missing detached child: %+v", result)
	}
	if actual := read_control(t, fmt.Sprintf("/proc/%d", pid), "cgroup"); actual != location {
		t.Fatalf("detached descendant escaped: %s", actual)
	}
	session.Disable()
	wait_signal(t, session.Done())
	if state, err := session.Status(); state != Stopped || err != nil {
		t.Fatalf("cleanup %s: %v", state, err)
	}
	if _, err := os.Stat(groupPath); !os.IsNotExist(err) {
		t.Fatalf("session group survived Done: %v", err)
	}
	// Kernel exit can remove a task from its cgroup just before /proc shows Z.
	wait_control(t, func() bool {
		data, err := os.ReadFile(fmt.Sprintf("/proc/%d/stat", pid))
		if os.IsNotExist(err) {
			return true
		}
		if err != nil {
			t.Fatal(err)
		}
		_, tail, found := strings.Cut(string(data), ") ")
		return found && strings.HasPrefix(tail, "Z ")
	})
	// A new identity can reuse capacity after all descendants are gone.
	resource_peer(t, manager, "detaching")
}

func TestNativeContainmentExec(t *testing.T) {
	manager := contained_manager(t, DefaultResourceLimits())
	group, err := create_cgroup(filepath.Join(manager.containment.group.path, "exec-proof"), DefaultResourceLimits().PerSession, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		if err := group.kill_and_remove(time.Second); err != nil {
			t.Error(err)
		}
	}()
	python, err := exec.LookPath("python3")
	if err != nil {
		t.Fatal(err)
	}
	fixture, err := filepath.Abs("testdata/resource_peer.py")
	if err != nil {
		t.Fatal(err)
	}
	command := exec.Command(python, "-I", "-u", fixture)
	command.SysProcAttr = &syscall.SysProcAttr{UseCgroupFD: true, CgroupFD: int(group.file.Fd())}
	command.Stdin = strings.NewReader(`{"id":"1","method":"initialize","params":{"protocol":{"major":1,"minor":0},"session":"test"}}` + "\n" + `{"method":"shutdown"}` + "\n")
	output, err := command.CombinedOutput()
	if err != nil || !strings.Contains(string(output), `"result"`) {
		t.Fatalf("private fixture exec: %v: %s", err, output)
	}
}

func TestNativeContainmentTaskLimits(t *testing.T) {
	resources := DefaultResourceLimits()
	resources.PerSession.Tasks, resources.Total.Tasks = 8, 12
	manager := contained_manager(t, resources)
	first := resource_peer(t, manager, "first")
	second := resource_peer(t, manager, "second")
	for _, session := range []*Session{first, second} {
		if result := resource_call(t, session, `{"action":"tasks"}`); result.Err != nil {
			t.Fatal(result.Err)
		}
	}
	parent := manager.containment.group.path
	if count := read_control(t, parent, "pids.current"); count != "12" {
		t.Fatalf("aggregate task limit was not reached: %s", count)
	}
	if count := read_control(t, filepath.Join(parent, first.Identity()), "pids.current"); count != "8" {
		t.Fatalf("session task limit was not reached: %s", count)
	}
	if control_counter(t, parent, "pids.events", "max") == 0 {
		t.Fatal("kernel did not report task limit enforcement")
	}
}

func TestNativeContainmentMemoryLimits(t *testing.T) {
	for _, aggregate := range []bool{false, true} {
		t.Run(fmt.Sprintf("aggregate-%t", aggregate), func(t *testing.T) {
			resources := DefaultResourceLimits()
			resources.PerSession.MemoryBytes, resources.Total.MemoryBytes = 64<<20, 96<<20
			manager := contained_manager(t, resources)
			first := resource_peer(t, manager, "first")
			if aggregate {
				second := resource_peer(t, manager, "second")
				if result := resource_call(t, first, `{"action":"memory","bytes":50331648}`); result.Err != nil {
					t.Fatal(result.Err)
				}
				// The kernel may select either session as the aggregate OOM victim.
				resource_call(t, second, `{"action":"memory","bytes":50331648}`)
			} else {
				if result := resource_call(t, first, `{"action":"memory","bytes":100663296}`); result.Err == nil {
					t.Fatal("allocation above session limit succeeded")
				}
				wait_signal(t, first.Done())
			}
			wait_control(t, func() bool {
				return control_counter(t, manager.containment.group.path, "memory.events", "oom_kill") > 0
			})
		})
	}
}

func TestNativeContainmentCPUThrottling(t *testing.T) {
	for _, aggregate := range []bool{false, true} {
		t.Run(fmt.Sprintf("aggregate-%t", aggregate), func(t *testing.T) {
			resources := DefaultResourceLimits()
			if aggregate {
				resources.PerSession.CPUQuota, resources.Total.CPUQuota = 50000, 50000
			}
			manager := contained_manager(t, resources)
			first := resource_peer(t, manager, "first")
			path := filepath.Join(manager.containment.group.path, first.Identity())
			if aggregate {
				second := resource_peer(t, manager, "second")
				if result := resource_call(t, second, `{"action":"cpu"}`); result.Err != nil {
					t.Fatal(result.Err)
				}
				path = manager.containment.group.path
			}
			if result := resource_call(t, first, `{"action":"cpu"}`); result.Err != nil {
				t.Fatal(result.Err)
			}
			wait_control(t, func() bool { return control_counter(t, path, "cpu.stat", "nr_throttled") > 0 })
		})
	}
}

func TestNativeContainmentFailedLaunchReleasesGroups(t *testing.T) {
	manager := contained_manager(t, DefaultResourceLimits())
	session, err := manager.Start(context.Background(), Spec{ID: "missing", Program: []string{"/nonexistent/tentacle-fixture"}})
	if err != nil {
		t.Fatal(err)
	}
	wait_signal(t, session.Done())
	if state, err := session.Status(); state != Failed || err == nil {
		t.Fatalf("missing executable %s: %v", state, err)
	}
	if _, err := os.Stat(filepath.Join(manager.containment.group.path, session.Identity())); !os.IsNotExist(err) {
		t.Fatalf("failed launch leaked a group: %v", err)
	}
	resource_peer(t, manager, "missing")
}

func TestNativeContainmentSDKRoundTrip(t *testing.T) {
	for _, language := range []string{"go", "python"} {
		t.Run(language, func(t *testing.T) {
			manager := contained_manager(t, DefaultResourceLimits())
			program := sdk_program(t, language, filepath.Join(t.TempDir(), "cancelled"))
			session, err := manager.Start(context.Background(), Spec{ID: "sdk", Program: program})
			if err != nil {
				t.Fatal(err)
			}
			wait_signal(t, session.Ready())
			if state, err := session.Status(); state != Ready || err != nil {
				t.Fatalf("contained SDK startup %s: %v", state, err)
			}
			if err := session.Activate(); err != nil {
				t.Fatal(err)
			}
			const arguments = `{"text":"é猫😀é"}`
			if result := resource_call(t, session, arguments); result.Err != nil || string(result.Value) != arguments {
				t.Fatalf("contained SDK round trip: %+v", result)
			}
		})
	}
}

func TestNativeContainmentCleanupFailure(t *testing.T) {
	manager := contained_manager(t, DefaultResourceLimits())
	session := resource_peer(t, manager, "cleanup-failure")
	groupPath := filepath.Join(manager.containment.group.path, session.Identity())
	killPath := filepath.Join(groupPath, "cgroup.kill")
	info, err := os.Stat(killPath)
	if err != nil {
		t.Fatal(err)
	}
	// This peer has no descendants. Remove kill-file access to exercise the
	// real cleanup error; direct child termination still joins the pipe workers.
	if err := os.Chmod(killPath, 0); err != nil {
		t.Fatal(err)
	}
	defer func() {
		if err := os.Chmod(killPath, info.Mode().Perm()); err != nil {
			t.Error(err)
			return
		}
		if err := (&cgroup{path: groupPath}).kill_and_remove(time.Second); err != nil {
			t.Error(err)
			return
		}
		// Only the fixture repairs its deliberately broken group; production
		// requires supervisor recovery and never reclaims an uncertain slot.
		manager.mutex.Lock()
		delete(manager.sessions, session.spec.ID)
		manager.cleanupError = nil
		manager.mutex.Unlock()
	}()
	session.Disable()
	wait_signal(t, session.Done())
	if state, err := session.Status(); state != Failed || !errors.Is(err, ErrContainment) {
		t.Fatalf("failed cleanup reported %s: %v", state, err)
	}
	if manager.Stats().Processes != 1 {
		t.Fatal("failed cleanup released admission capacity")
	}
	if _, err := manager.Start(context.Background(), child_spec(t, "blocked", "echo")); !errors.Is(err, ErrContainment) {
		t.Fatalf("cleanup failure admitted a child: %v", err)
	}
	if err := manager.Shutdown(context.Background()); !errors.Is(err, ErrContainment) {
		t.Fatalf("cleanup failure reported successful shutdown: %v", err)
	}
}