package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
)

// https://kernel.googlesource.com/pub/scm/linux/kernel/git/torvalds/linux/+/v4.17/include/linux/sched/topology.h#20
var SDFlags = map[string]uint64{
	"SD_LOAD_BALANCE":        0x0001,
	"SD_BALANCE_NEWIDLE":     0x0002,
	"SD_BALANCE_EXEC":        0x0004,
	"SD_BALANCE_FORK":        0x0008,
	"SD_BALANCE_WAKE":        0x0010,
	"SD_WAKE_AFFINE":         0x0020,
	"SD_ASYM_CPUCAPACITY":    0x0040,
	"SD_SHARE_CPUCAPACITY":   0x0080,
	"SD_SHARE_POWERDOMAIN":   0x0100,
	"SD_SHARE_PKG_RESOURCES": 0x0200,
	"SD_SERIALIZE":           0x0400,
	"SD_ASYM_PACKING":        0x0800,
	"SD_PREFER_SIBLING":      0x1000,
	"SD_OVERLAP":             0x2000,
	"SD_NUMA":                0x4000,
}

type Scheduler map[string][]Domain

type Domain struct {
	Name    string            `json:"name"`
	Type    string            `json:"type"`
	Flags   []string          `json:"flags"`
	Indexes map[string]string `json:"indexes"`
}

func main() {
	cpus, err := CPUs()
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		os.Exit(1)
	}

	if len(cpus) == 0 {
		fmt.Fprint(os.Stderr, "there is no scheduler domains\n")
		os.Exit(1)
	}

	sched := Scheduler{}
	for _, cpu := range cpus {
		_, cpuID := path.Split(cpu)
		domains, err := domains(cpu)
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			os.Exit(1)
		}
		sched[cpuID] = domains
	}
	out, err := json.Marshal(sched)
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		os.Exit(1)
	}
	fmt.Println(string(out))
}

func domains(cpuPath string) ([]Domain, error) {
	domainPath := fmt.Sprintf("%s/domain*", cpuPath)
	availDomains, err := filepath.Glob(domainPath)
	if err != nil {
		return nil, fmt.Errorf("failed to get domains under %s: %v", cpuPath, err)
	}

	domains := []Domain{}

	if len(availDomains) == 0 {
		return domains, nil
	}

	for _, d := range availDomains {
		_, dName := path.Split(d)
		dType := getContent(d, "name")
		flags, err := domainFlags(d)
		if err != nil {
			return nil, err
		}
		indexes := domainIndexes(d)

		domain := Domain{
			Name:    dName,
			Type:    dType,
			Flags:   flags,
			Indexes: indexes,
		}
		domains = append(domains, domain)
	}
	return domains, nil
}

func domainFlags(path string) ([]string, error) {
	flagPath := fmt.Sprintf("%s/flags", path)

	content, err := ioutil.ReadFile(flagPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read %s: %v", flagPath, err)
	}

	flags, err := strconv.ParseUint(strings.TrimSpace(string(content)), 0, 64)
	if err != nil {
		return nil, fmt.Errorf("failed to convert flags %s: %v", flagPath, err)
	}

	supportedFlags := []string{}
	for k, v := range SDFlags {
		if flags&v > 0 {
			supportedFlags = append(supportedFlags, k)
		}
	}
	return supportedFlags, nil
}

func domainIndexes(path string) map[string]string {
	indexes := map[string]string{
		"busy":      getContent(path, "busy_idx"),
		"idle":      getContent(path, "idle_idx"),
		"new_idle":  getContent(path, "newidle_idx"),
		"wake":      getContent(path, "wake_idx"),
		"fork_exec": getContent(path, "forkexec_idx"),
	}
	return indexes
}

func getContent(path, fileName string) string {
	domainName := fmt.Sprintf("%s/%s", path, fileName)
	name, err := ioutil.ReadFile(domainName)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(name))
}

func CPUs() ([]string, error) {
	cpus, err := filepath.Glob("/proc/sys/kernel/sched_domain/cpu*")
	if err != nil {
		return nil, fmt.Errorf("failed to get a list of cpus: %v", err)
	}
	return cpus, nil
}