package memlimit

import (
	"os"
	"path/filepath"
	"runtime/debug"
	"testing"
)

// --- readLimit unit tests ---

func TestReadLimit(t *testing.T) {
	tests := []struct {
		name    string
		content string
		want    int64
		ok      bool
	}{
		{"normal", "1073741824\n", 1073741824, true},
		{"max_keyword", "max\n", 0, false},
		{"empty", "", 0, false},
		{"whitespace_only", "  \n", 0, false},
		{"unlimited_v1", "9223372036854771712\n", 0, false},
		{"negative", "-1\n", 0, false},
		{"zero", "0\n", 0, false},
		{"512Gi_plausible", "549755813888\n", 549755813888, true},
		{"just_over_1TiB", "1099511627777\n", 0, false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			path := filepath.Join(t.TempDir(), "memory.max")
			if err := os.WriteFile(path, []byte(tt.content), 0o644); err != nil {
				t.Fatal(err)
			}
			got, ok := readLimit(path)
			if ok != tt.ok {
				t.Fatalf("ok = %v, want %v", ok, tt.ok)
			}
			if ok && got != tt.want {
				t.Fatalf("got %d, want %d", got, tt.want)
			}
		})
	}
}

func TestReadLimit_FileNotExists(t *testing.T) {
	_, ok := readLimit("/nonexistent/path/memory.max")
	if ok {
		t.Fatal("expected ok=false for nonexistent file")
	}
}

// --- detect tests (with cgroup path stubbing) ---

func swapCgroupRoot(t *testing.T, tmp string) func() {
	t.Helper()
	origV2 := detectV2Path
	origV1 := detectV1Path
	detectV2Path = filepath.Join(tmp, "memory.max")
	detectV1Path = filepath.Join(tmp, "memory", "memory.limit_in_bytes")
	return func() {
		detectV2Path = origV2
		detectV1Path = origV1
	}
}

func writeCgroupFile(t *testing.T, path, value string) {
	t.Helper()
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, []byte(value), 0o644); err != nil {
		t.Fatal(err)
	}
}

func TestDetect_CgroupV2(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()
	writeCgroupFile(t, detectV2Path, "1048576000") // 1000 MiB

	limit, source, ok := detect()
	if !ok || source != "cgroup.v2" || limit != 1048576000 {
		t.Fatalf("detect() = %d, %s, %v; want 1048576000, cgroup.v2, true", limit, source, ok)
	}
}

func TestDetect_CgroupV1(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()
	writeCgroupFile(t, detectV1Path, "2147483648") // 2 GiB

	limit, source, ok := detect()
	if !ok || source != "cgroup.v1" || limit != 2147483648 {
		t.Fatalf("detect() = %d, %s, %v; want 2147483648, cgroup.v1, true", limit, source, ok)
	}
}

func TestDetect_V2TakesPrecedence(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()
	writeCgroupFile(t, detectV2Path, "1000")
	writeCgroupFile(t, detectV1Path, "2000")

	limit, source, ok := detect()
	if !ok || source != "cgroup.v2" || limit != 1000 {
		t.Fatalf("v2 should take precedence: %d, %s, %v", limit, source, ok)
	}
}

func TestDetect_NoCgroup(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()

	_, _, ok := detect()
	if ok {
		t.Fatal("expected ok=false when no cgroup files exist")
	}
}

// --- SetupFromCgroup integration tests ---

// currentMemLimit reads the current soft memory limit by setting a throwaway
// value and immediately restoring it.
func currentMemLimit() int64 {
	current := debug.SetMemoryLimit(0) // 0 disables, returns previous
	debug.SetMemoryLimit(current)
	return current
}

func TestSetupFromCgroup_RespectsEnvVar(t *testing.T) {
	t.Setenv("GOMEMLIMIT", "500MiB")
	defer swapCgroupRoot(t, t.TempDir())()
	writeCgroupFile(t, detectV2Path, "1048576000")

	sentinel := int64(777777)
	original := debug.SetMemoryLimit(sentinel)
	defer debug.SetMemoryLimit(original)

	SetupFromCgroup()

	if got := currentMemLimit(); got != sentinel {
		t.Fatalf("GOMEMLIMIT env var should prevent override: want %d, got %d", sentinel, got)
	}
}

func TestSetupFromCgroup_Sets90Percent(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()
	writeCgroupFile(t, detectV2Path, "1048576000") // 1000 MiB

	original := debug.SetMemoryLimit(0)
	defer debug.SetMemoryLimit(original)

	SetupFromCgroup()

	expected := int64(943718400) // 1048576000 * 0.9
	if got := currentMemLimit(); got != expected {
		t.Fatalf("limit = %d, want %d", got, expected)
	}
}

func TestSetupFromCgroup_NoCgroup_NoOp(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()

	original := debug.SetMemoryLimit(0) // disable = effectively "no limit"
	defer debug.SetMemoryLimit(original)

	SetupFromCgroup()

	if got := currentMemLimit(); got != 0 {
		t.Fatalf("expected no limit change when cgroup absent, got %d", got)
	}
}

func TestSetupFromCgroup_ClampsToFloor(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()
	// 950 MiB cgroup → 90% = 855 MiB → below 900 MiB floor
	// cgroup limit (950 MiB) > floor (900 MiB), so clamping is safe
	writeCgroupFile(t, detectV2Path, "996147200") // 950 MiB

	original := debug.SetMemoryLimit(0)
	defer debug.SetMemoryLimit(original)

	SetupFromCgroup()

	expected := minGoMemLimit
	if got := currentMemLimit(); got != expected {
		t.Fatalf("limit = %d, want floor %d", got, expected)
	}
}

func TestSetupFromCgroup_FloorNotAppliedWhenCgroupBelowFloor(t *testing.T) {
	defer swapCgroupRoot(t, t.TempDir())()
	// 800 MiB cgroup → 90% = 720 MiB
	// cgroup limit (800 MiB) < floor (900 MiB), so floor must NOT be
	// applied — otherwise GOMEMLIMIT would exceed the container limit.
	writeCgroupFile(t, detectV2Path, "838860800") // 800 MiB

	original := debug.SetMemoryLimit(0)
	defer debug.SetMemoryLimit(original)

	SetupFromCgroup()

	expected := int64(754974720) // 838860800 * 0.9
	if got := currentMemLimit(); got != expected {
		t.Fatalf("limit = %d, want %d (floor must not exceed cgroup limit)", got, expected)
	}
}