package main

import (
	"io/fs"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

// cgroupIO is a cgroup's bytes and requests to storage, from its io.stat. Unlike /proc/<pid>/io, io.stat
// is readable by every account, so a unit or a container is seen whoever runs the tool runner.
type cgroupIO struct {
	read, write, rios, wios uint64
}

// parseIOStat sums the lines of an io.stat ("259:0 rbytes=… wbytes=… rios=… wios=…") over the devices
// counted counts; every device when counted is nil.
func parseIOStat(s string, counted map[string]bool) cgroupIO {
	var c cgroupIO
	for _, line := range strings.Split(s, "\n") {
		f := strings.Fields(line)
		if len(f) < 2 || (counted != nil && !counted[f[0]]) {
			continue
		}
		for _, kv := range f[1:] {
			k, v, _ := strings.Cut(kv, "=")
			n, _ := strconv.ParseUint(v, 10, 64)
			switch k {
			case "rbytes":
				c.read += n
			case "wbytes":
				c.write += n
			case "rios":
				c.rios += n
			case "wios":
				c.wios += n
			}
		}
	}
	return c
}

// bottomDevices are the major:minor of the disks I/O finally lands on. A write through an encrypted
// or logical volume is charged to the volume and again to the disk under it, and a loop device's to the
// filesystem its file is on, so only a whole disk built on nothing is counted. nil (count every line)
// when /proc/diskstats cannot be read.
func (m *Machine) bottomDevices() map[string]bool {
	text := m.read("/proc/diskstats")
	if text == "" {
		return nil
	}
	stats := parseDiskstats(text)
	disks := map[string]bool{}
	for _, name := range m.blockDevices(stats, false) {
		if len(m.below(name)) == 0 {
			disks[name] = true
		}
	}
	out := map[string]bool{}
	for _, line := range strings.Split(text, "\n") {
		f := strings.Fields(line)
		if len(f) > 2 && disks[f[2]] {
			out[f[0]+":"+f[1]] = true
		}
	}
	return out
}

// unitPaths is every systemd unit's cgroup path under /sys/fs/cgroup: the innermost .service or
// .scope, because a cgroup's own stat file counts its descendants too and an outer unit would count
// them twice. Shared by every reader that walks the cgroup tree (unitIO here, unitCPU in cpu.go).
func (m *Machine) unitPaths() []string {
	base := m.path("/sys/fs/cgroup")
	var units []string
	filepath.WalkDir(base, func(p string, d fs.DirEntry, err error) error {
		if err != nil || !d.IsDir() {
			return nil
		}
		if n := d.Name(); strings.HasSuffix(n, ".service") || strings.HasSuffix(n, ".scope") {
			rel, _ := filepath.Rel(base, p)
			units = append(units, filepath.ToSlash(rel))
		}
		return nil
	})
	outer := map[string]bool{}
	for _, u := range units {
		for dir := filepath.Dir(u); dir != "." && dir != "/"; dir = filepath.Dir(dir) {
			outer[dir] = true
		}
	}
	out := make([]string, 0, len(units))
	for _, u := range units {
		if !outer[u] {
			out = append(out, u)
		}
	}
	return out
}

// unitIO is every unit's storage I/O, by its cgroup's path.
func (m *Machine) unitIO(counted map[string]bool) map[string]cgroupIO {
	out := map[string]cgroupIO{}
	for _, u := range m.unitPaths() {
		text := m.read("/sys/fs/cgroup/" + u + "/io.stat")
		if text == "" {
			continue
		}
		out[u] = parseIOStat(text, counted)
	}
	return out
}

// topUnits is the units, or the containers, that moved the most bytes between two readings of io.stat.
func (m *Machine) topUnits(a, b map[string]cgroupIO, window time.Duration, by string, limit int) []Process {
	s := window.Seconds()
	groups := map[string]*Process{}
	var names map[string]string
	for path, end := range b {
		start, ok := a[path]
		if !ok {
			continue
		}
		r, w := sub(end.read, start.read), sub(end.write, start.write)
		if r+w == 0 {
			continue
		}
		unit := filepath.Base(path)
		key := unit
		var container string
		if id := containerID.FindStringSubmatch("/" + path); id != nil {
			if names == nil {
				names = m.containerNames()
			}
			container = id[1][:12]
			if n, ok := names[id[1]]; ok {
				container = n
			}
		}
		if by == "container" {
			if container == "" {
				continue
			}
			key = container
		}
		g, ok := groups[key]
		if !ok {
			g = &Process{Unit: unit, Container: container}
			if by == "container" {
				g.Unit = ""
			}
			groups[key] = g
		}
		g.ReadMiB += r
		g.WriteMiB += w
		g.bytes += r + w
		g.ReadIOPS += round1(sub(end.rios, start.rios) / s)
		g.WriteIOPS += round1(sub(end.wios, start.wios) / s)
	}
	out := make([]Process, 0, len(groups))
	for _, g := range groups {
		out = append(out, *g)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].bytes != out[j].bytes {
			return out[i].bytes > out[j].bytes
		}
		return out[i].Unit+out[i].Container < out[j].Unit+out[j].Container
	})
	if len(out) > limit {
		out = out[:limit]
	}
	for i := range out {
		inMiB(&out[i], s)
	}
	return out
}