package main

import (
	"fmt"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
)

// THPDir is where the kernel says and takes its transparent huge pages.
const THPDir = "sys/kernel/mm/transparent_hugepage"

// CgroupRoot is the unified cgroup hierarchy.
const CgroupRoot = "sys/fs/cgroup"

// HugePagesAnswer is what kernel_huge_pages answers.
type HugePagesAnswer struct {
	// Mode is the word the kernel brackets in enabled: always, madvise or never.
	Mode    string `json:"mode"`
	Defrag  string `json:"defrag,omitempty"`
	Shmem   string `json:"shmem_enabled,omitempty"`
	PerSize []Size `json:"per_size,omitempty"`
	// Meminfo is the machine's huge pages by kind, in MiB.
	Meminfo map[string]float64 `json:"meminfo_mib,omitempty"`
	Units   []Unit             `json:"units"`
	Said    []string           `json:"said"`
}

// Size is one per-size control (hugepages-<n>kB/enabled).
type Size struct {
	Size    string `json:"size"`
	Enabled string `json:"enabled"`
	// Means is what inherit resolves to: the top-level mode.
	Means string `json:"means"`
}

// Unit is one cgroup holding file pages in huge folios.
type Unit struct {
	Cgroup          string `json:"cgroup"`
	FileTHP         int64  `json:"file_thp_bytes"`
	FileDirty       int64  `json:"file_dirty_bytes"`
	FileWriteback   int64  `json:"file_writeback_bytes"`
	FileMapped      int64  `json:"file_mapped_bytes"`
	CollapsedByKhug int64  `json:"thp_collapse_alloc"`
	Written         int64  `json:"written_bytes"`
}

// bracketed is the word a sysfs choice file brackets: "always [madvise] never" is madvise.
func bracketed(text string) string {
	open := strings.IndexByte(text, '[')
	end := strings.IndexByte(text, ']')
	if open < 0 || end < open {
		return ""
	}
	return text[open+1 : end]
}

// HugePages reads the mode, the per-size controls, meminfo and the top units by file_thp.
func (m *Machine) HugePages(units int) (*HugePagesAnswer, error) {
	enabled, err := m.read(THPDir + "/enabled")
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fmt.Errorf("this kernel has no /%s/enabled: it was built without transparent huge pages", THPDir)
		}
		return nil, fmt.Errorf("reading /%s/enabled: %w", THPDir, err)
	}
	a := &HugePagesAnswer{Mode: bracketed(enabled), Units: []Unit{}}
	if s, err := m.read(THPDir + "/defrag"); err == nil {
		a.Defrag = bracketed(s)
	}
	if s, err := m.read(THPDir + "/shmem_enabled"); err == nil {
		a.Shmem = bracketed(s)
	}
	a.Said = append(a.Said, fmt.Sprintf("transparent huge pages: %s, defrag %s", a.Mode, orUnknown(a.Defrag)))

	sizes, _ := m.dirs(THPDir)
	for _, d := range sizes {
		if !strings.HasPrefix(d, "hugepages-") {
			continue
		}
		s, err := m.read(THPDir + "/" + d + "/enabled")
		if err != nil {
			continue
		}
		sz := Size{Size: strings.TrimPrefix(d, "hugepages-"), Enabled: bracketed(s)}
		sz.Means = sz.Enabled
		if sz.Enabled == "inherit" {
			sz.Means = a.Mode
		}
		a.PerSize = append(a.PerSize, sz)
	}
	sort.Slice(a.PerSize, func(i, j int) bool { return kib(a.PerSize[i].Size) < kib(a.PerSize[j].Size) })
	var on []string
	for _, s := range a.PerSize {
		if s.Means != "never" {
			on = append(on, s.Size+" "+s.Means)
		}
	}
	if len(a.PerSize) > 0 {
		a.Said = append(a.Said, "sizes not never: "+orNone(on))
	}

	if mi, err := m.read("proc/meminfo"); err == nil {
		a.Meminfo = meminfo(mi)
		a.Said = append(a.Said, fmt.Sprintf("huge pages held: anonymous %.0f MiB, shared memory %.0f MiB, file %.0f MiB (mapped at 2 MiB %.0f MiB)",
			a.Meminfo["AnonHugePages"], a.Meminfo["ShmemHugePages"], a.Meminfo["FileHugePages"], a.Meminfo["FilePmdMapped"]))
	}

	if units > 0 {
		all := m.cgroups()
		sort.Slice(all, func(i, j int) bool {
			if all[i].FileTHP != all[j].FileTHP {
				return all[i].FileTHP > all[j].FileTHP
			}
			return all[i].Cgroup < all[j].Cgroup
		})
		if len(all) > units {
			all = all[:units]
		}
		a.Units = all
		for _, u := range all {
			a.Said = append(a.Said, fmt.Sprintf("%s: file in huge folios %s, dirty %s, writeback %s, written %s since it started",
				u.Cgroup, mib(u.FileTHP), mib(u.FileDirty), mib(u.FileWriteback), mib(u.Written)))
		}
		if len(all) == 0 {
			a.Said = append(a.Said, "no unit holds file pages in huge folios")
		}
	}
	return a, nil
}

// cgroups walks the leaf services and scopes (to three levels below the root) that hold file_thp.
func (m *Machine) cgroups() []Unit {
	var out []Unit
	var walk func(rel string, depth int)
	walk = func(rel string, depth int) {
		children, _ := m.dirs(path.Join(CgroupRoot, rel))
		for _, c := range children {
			p := path.Join(rel, c)
			if strings.HasSuffix(c, ".slice") {
				if depth < 3 {
					walk(p, depth+1)
				}
				continue
			}
			if !strings.HasSuffix(c, ".service") && !strings.HasSuffix(c, ".scope") {
				continue
			}
			stat, err := m.read(path.Join(CgroupRoot, p, "memory.stat"))
			if err != nil {
				continue
			}
			kv := memoryStat(stat)
			if kv["file_thp"] == 0 {
				continue
			}
			u := Unit{Cgroup: p, FileTHP: kv["file_thp"], FileDirty: kv["file_dirty"], FileWriteback: kv["file_writeback"],
				FileMapped: kv["file_mapped"], CollapsedByKhug: kv["thp_collapse_alloc"]}
			if io, err := m.read(path.Join(CgroupRoot, p, "io.stat")); err == nil {
				u.Written = mostWritten(io)
			}
			out = append(out, u)
		}
	}
	walk("", 0)
	return out
}

// memoryStat reads memory.stat's "key value" lines.
func memoryStat(text string) map[string]int64 {
	kv := map[string]int64{}
	for _, l := range strings.Split(text, "\n") {
		f := strings.Fields(l)
		if len(f) != 2 {
			continue
		}
		if n, err := strconv.ParseInt(f[1], 10, 64); err == nil {
			kv[f[0]] = n
		}
	}
	return kv
}

// mostWritten is the most bytes any one device of io.stat saw written: a mirror's member disks and the
// mirror itself each count the same write, so the lines are never summed.
func mostWritten(text string) int64 {
	var most int64
	for _, l := range strings.Split(text, "\n") {
		for _, f := range strings.Fields(l) {
			if v, ok := strings.CutPrefix(f, "wbytes="); ok {
				if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > most {
					most = n
				}
			}
		}
	}
	return most
}

// meminfo reads the huge page lines of /proc/meminfo, in MiB.
func meminfo(text string) map[string]float64 {
	want := map[string]bool{"AnonHugePages": true, "ShmemHugePages": true, "ShmemPmdMapped": true, "FileHugePages": true, "FilePmdMapped": true}
	out := map[string]float64{}
	for _, l := range strings.Split(text, "\n") {
		k, v, ok := strings.Cut(l, ":")
		if !ok || !want[k] {
			continue
		}
		f := strings.Fields(v)
		if len(f) == 0 {
			continue
		}
		if n, err := strconv.ParseFloat(f[0], 64); err == nil {
			out[k] = n / 1024
		}
	}
	return out
}

func kib(size string) int {
	n, _ := strconv.Atoi(strings.TrimSuffix(size, "kB"))
	return n
}

func mib(b int64) string { return fmt.Sprintf("%.1f MiB", float64(b)/(1<<20)) }

func orUnknown(s string) string {
	if s == "" {
		return "unknown"
	}
	return s
}

func orNone(list []string) string {
	if len(list) == 0 {
		return "none"
	}
	return strings.Join(list, ", ")
}