package budget

import (
	"bytes"
	"math"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"

	"github.com/BurntSushi/toml"
	"github.com/godbus/dbus/v5"
)

func testConfig() config {
	return config{
		Profiles: map[string]profile{
			"default":      {"400%", "6G", "8G", "2G"},
			"compute":      {"800%", "12G", "16G", "2G"},
			"bulk":         {"400%", "6G", "8G", "2G"},
			"search":       {"200%", "512M", "1G", "0"},
			"unrestricted": {"", "infinity", "infinity", "infinity"},
		},
		Devices: []device{{"/dev/mock", map[string]ioLimit{
			"default": {"100M", "100M", "4000", "4000"},
			"compute": {"500M", "500M", "20000", "20000"},
			"bulk":    {"500M", "500M", "10000", "10000"},
			"search":  {"50M", "5M", "1000", "1000"},
		}}},
		Apps: map[string][]string{
			"unrestricted": {"org.example.Terminal", "org.example.Console"},
			"compute":      {"game", "org.example.Game"},
			"bulk":         {"org.example.Download"},
		},
		Units:     map[string][]string{"unrestricted": {"containerd.service", "worker.service"}},
		Prefixes:  map[string][]string{"unrestricted": {"app-terminal-.scope.d"}},
		Commands:  map[string]string{"rg": "search", "java": "compute", "rsync": "bulk"},
		Python:    "compute",
		Terminals: []string{"app-terminal-", "app-org.example.Terminal", "app-org.example.Console"},
		Shared:    map[string]string{"editor": "default", "editor-uri-handler": "default"},
	}
}

func values(properties []property) map[string]dbus.Variant {
	result := map[string]dbus.Variant{}
	for _, p := range properties {
		result[p.Name] = p.Value
	}
	return result
}

func TestParseSize(t *testing.T) {
	for _, tc := range []struct {
		value string
		base  uint64
		want  uint64
	}{
		{"8G", 1024, 8589934592}, {"12G", 1024, 12884901888},
		{"16G", 1024, 17179869184}, {"512M", 1024, 536870912},
		{"50M", 1000, 50000000}, {"5K", 1000, 5000},
		{"1000", 1000, 1000}, {"0", 1024, 0}, {"infinity", 1024, math.MaxUint64},
	} {
		t.Run(tc.value, func(t *testing.T) {
			got, err := parseSize(tc.value, tc.base)
			if err != nil || got != tc.want {
				t.Fatalf("got %d, %v; want %d", got, err, tc.want)
			}
		})
	}
	for _, input := range []string{"", "-1", "abc", "1T"} {
		if _, err := parseSize(input, 1024); err == nil {
			t.Errorf("accepted %q", input)
		}
	}
}

func TestComputeProps(t *testing.T) {
	props, err := testConfig().properties("compute")
	if err != nil {
		t.Fatal(err)
	}
	got := values(props)
	for name, want := range map[string]uint64{
		"CPUQuotaPerSecUSec": 8000000, "MemoryHigh": 12884901888,
		"MemoryMax": 17179869184, "MemorySwapMax": 2147483648,
	} {
		if got[name].Signature().String() != "t" || got[name].Value() != want {
			t.Errorf("%s: %v", name, got[name])
		}
	}
	if got["IOReadBandwidthMax"].Signature().String() != "a(st)" {
		t.Fatal("wrong I/O wire signature")
	}
	want := []diskLimit{{"/dev/mock", 500000000}}
	if !reflect.DeepEqual(got["IOReadBandwidthMax"].Value(), want) {
		t.Fatal(got["IOReadBandwidthMax"])
	}
	if dbus.SignatureOf(props).String() != "a(sv)" {
		t.Fatal("wrong property wire signature")
	}
	if dbus.SignatureOf([]auxiliary{}).String() != "a(sa(sv))" {
		t.Fatal("wrong auxiliary wire signature")
	}
}

func TestUnrestricted(t *testing.T) {
	c := testConfig()
	props, err := c.properties("unrestricted")
	if err != nil {
		t.Fatal(err)
	}
	got := values(props)
	for _, name := range []string{"CPUQuotaPerSecUSec", "MemoryHigh", "MemoryMax", "MemorySwapMax"} {
		if got[name].Value() != uint64(math.MaxUint64) {
			t.Errorf("%s wasn't cleared", name)
		}
	}
	for _, field := range ioFields {
		if got[field.name].Signature().String() != "a(st)" || len(got[field.name].Value().([]diskLimit)) != 0 {
			t.Errorf("%s wasn't cleared", field.name)
		}
	}
	assignments, err := c.assignments("unrestricted")
	if err != nil {
		t.Fatal(err)
	}
	for _, want := range []string{"CPUQuota=", "MemoryMax=infinity", "IOReadBandwidthMax=", "IOWriteIOPSMax="} {
		if !strings.Contains(strings.Join(assignments, "\n"), want+"\n") && assignments[len(assignments)-1] != want {
			t.Errorf("missing reset %q", want)
		}
	}
}

func TestProfileErrors(t *testing.T) {
	c := testConfig()
	if _, err := c.properties("missing"); err == nil {
		t.Fatal("unknown profile accepted")
	}
	if _, err := c.assignments("missing"); err == nil {
		t.Fatal("unknown profile accepted")
	}
	for _, p := range []profile{{"bad%", "6G", "8G", "2G"}, {"400%", "bad", "8G", "2G"}} {
		c.Profiles["bad"] = p
		if _, err := c.properties("bad"); err == nil {
			t.Fatal("invalid size accepted")
		}
	}
	c.Devices[0].Limits["default"] = ioLimit{"bad", "100M", "4000", "4000"}
	if _, err := c.properties("default"); err == nil {
		t.Fatal("invalid disk rate accepted")
	}
}

func TestLoadConfig(t *testing.T) {
	path := filepath.Join(t.TempDir(), "profiles.toml")
	data := `python = "compute"
terminals = ["app-terminal-", "app-org.example.Console"]

[profiles.compute]
cpu = "800%"
high = "12G"
memory = "16G"
swap = "2G"

[[devices]]
path = '/dev/mock'
[devices.limits.compute]
read = "500M"
write = "500M"
riops = "20000"
wiops = "20000"

[apps]
compute = ["game", "org.example.Game"]
[units]
unrestricted = ["containerd.service"]
[prefixes]
unrestricted = ["app-terminal-.scope.d"]
[commands]
java = "compute"
[shared]
"editor-uri-handler" = "default"
`
	if err := os.WriteFile(path, []byte(data), 0600); err != nil {
		t.Fatal(err)
	}
	c, err := loadConfig(path)
	if err != nil {
		t.Fatal(err)
	}
	want := testConfig()
	want.Profiles = map[string]profile{"compute": want.Profiles["compute"]}
	want.Devices[0].Limits = map[string]ioLimit{"compute": want.Devices[0].Limits["compute"]}
	want.Apps = map[string][]string{"compute": want.Apps["compute"]}
	want.Units["unrestricted"] = []string{"containerd.service"}
	want.Commands = map[string]string{"java": "compute"}
	want.Terminals = []string{"app-terminal-", "app-org.example.Console"}
	want.Shared = map[string]string{"editor-uri-handler": "default"}
	if !reflect.DeepEqual(c, want) {
		t.Fatalf("got %#v, want %#v", c, want)
	}
	if err := os.WriteFile(path, []byte("{"), 0600); err != nil {
		t.Fatal(err)
	}
	if _, err := loadConfig(path); err == nil {
		t.Fatal("invalid TOML accepted")
	}
	if _, err := loadConfig(path + "-missing"); !os.IsNotExist(err) {
		t.Fatal(err)
	}
}

func TestConfigRoundTrip(t *testing.T) {
	want := testConfig()
	var out bytes.Buffer
	if err := toml.NewEncoder(&out).Encode(want); err != nil {
		t.Fatal(err)
	}
	var got config
	if _, err := toml.Decode(out.String(), &got); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("got %#v, want %#v", got, want)
	}
}

func TestExample(t *testing.T) {
	c, err := loadConfig(filepath.Join("..", "..", "profiles.example.toml"))
	if err != nil {
		t.Fatal(err)
	}
	if _, err := c.render(); err != nil {
		t.Fatal(err)
	}
	for name := range c.Profiles {
		if _, err := c.properties(name); err != nil {
			t.Fatal(name, err)
		}
	}
}