package metrics

import (
	"bytes"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// A fake /sys/class/powercap: a package with its parts, the platform domain, and the MMIO
// copy of the same package. Only the package counts.
func fakeRAPL(t *testing.T, domains map[string][2]string) string {
	t.Helper()
	root := t.TempDir()
	for dir, v := range domains {
		d := filepath.Join(root, dir)
		if err := os.MkdirAll(d, 0o755); err != nil {
			t.Fatal(err)
		}
		os.WriteFile(filepath.Join(d, "name"), []byte(v[0]+"\n"), 0o644)
		os.WriteFile(filepath.Join(d, "energy_uj"), []byte(v[1]+"\n"), 0o644)
	}
	return root
}

func TestRAPLPackagesOnlyCountsPackages(t *testing.T) {
	root := fakeRAPL(t, map[string][2]string{
		"intel-rapl:0":      {"package-0", "1000"},
		"intel-rapl:0:0":    {"core", "600"},
		"intel-rapl:0:1":    {"uncore", "10"},
		"intel-rapl:1":      {"psys", "300"},
		"intel-rapl-mmio:0": {"package-0", "1000"},
	})
	if got := raplPackages(root); len(got) != 1 || got[0] != "intel-rapl:0" {
		t.Fatalf("packages = %v, want [intel-rapl:0]", got)
	}

	two := fakeRAPL(t, map[string][2]string{"intel-rapl:0": {"package-0", "5"}, "intel-rapl:1": {"package-1", "7"}})
	got := raplPackages(two)
	if strings.Join(got, ",") != "intel-rapl:0,intel-rapl:1" {
		t.Fatalf("two sockets: packages = %v", got)
	}
	cs, ok := readRAPL(two, got)
	if !ok || len(cs) != 2 || cs[0].UJ != 5 || cs[1].UJ != 7 {
		t.Fatalf("readRAPL = %+v, %v", cs, ok)
	}
	if got := raplPackages(filepath.Join(two, "nope")); len(got) != 0 {
		t.Fatalf("no powercap: packages = %v, want none", got)
	}
}

func TestRAPLWatts(t *testing.T) {
	c := func(d string, uj uint64) raplCount { return raplCount{d, uj} }
	max := map[string]uint64{"intel-rapl:0": 1_000_000_000}
	cases := []struct {
		name      string
		prev, cur []raplCount
		dt        float64
		want      float64
		ok        bool
	}{
		{"steady", []raplCount{c("intel-rapl:0", 0)}, []raplCount{c("intel-rapl:0", 10_000_000)}, 2, 5, true},
		{"two packages add", []raplCount{c("intel-rapl:0", 0), c("intel-rapl:1", 0)}, []raplCount{c("intel-rapl:0", 4_000_000), c("intel-rapl:1", 6_000_000)}, 1, 10, true},
		{"wrapped", []raplCount{c("intel-rapl:0", 999_000_000)}, []raplCount{c("intel-rapl:0", 1_000_000)}, 1, 2, true},
		{"backwards, range unknown", []raplCount{c("intel-rapl:1", 5_000_000)}, []raplCount{c("intel-rapl:1", 1)}, 1, 0, false},
		{"counters changed", []raplCount{c("intel-rapl:0", 0)}, []raplCount{c("intel-rapl:1", 9)}, 1, 0, false},
		{"first reading", nil, []raplCount{c("intel-rapl:0", 9)}, 1, 0, false},
		{"no time passed", []raplCount{c("intel-rapl:0", 0)}, []raplCount{c("intel-rapl:0", 9)}, 0, 0, false},
	}
	for _, tc := range cases {
		got, ok := raplWatts(tc.prev, tc.cur, max, tc.dt)
		if ok != tc.ok || (ok && (got < tc.want-1e-9 || got > tc.want+1e-9)) {
			t.Errorf("%s: got %v, %v; want %v, %v", tc.name, got, ok, tc.want, tc.ok)
		}
	}
}

// The helper answers every newline with the counters it found itself, and stops at EOF.
func TestCPUEnergyHelper(t *testing.T) {
	root := fakeRAPL(t, map[string][2]string{"intel-rapl:0": {"package-0", "42"}, "intel-rapl:0:0": {"core", "7"}})
	var out bytes.Buffer
	if code := cpuEnergyHelper(root, strings.NewReader("\n\n"), &out); code != 0 {
		t.Fatalf("exit %d", code)
	}
	lines := strings.Split(strings.TrimSpace(out.String()), "\n")
	if len(lines) != 2 {
		t.Fatalf("got %d lines, want 2: %q", len(lines), out.String())
	}
	var cs []raplCount
	if err := json.Unmarshal([]byte(lines[1]), &cs); err != nil || len(cs) != 1 || cs[0] != (raplCount{"intel-rapl:0", 42}) {
		t.Fatalf("line = %s (%v)", lines[1], err)
	}
}

// History keeps the average of the samples that had a reading.
func TestFoldCPUPower(t *testing.T) {
	b := []Sample{{T: 1, CPUPower: 4}, {T: 2}, {T: 3, CPUPower: 8}}
	if got := fold(b).CPUPower; got != 6 {
		t.Fatalf("folded CPU power = %v, want 6", got)
	}
	if got := fold([]Sample{{T: 1}, {T: 2}}).CPUPower; got != 0 {
		t.Fatalf("no readings folded to %v, want 0", got)
	}
}