package main

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"
)

// cpuTimes is the first line of /proc/stat, in clock ticks: what every CPU together spent its time on.
type cpuTimes struct {
	user, nice, system, idle, iowait, irq, softirq, steal uint64
	cpus                                                  int
}

// total leaves guest time out: the kernel already counts it in user and nice.
func (c cpuTimes) total() uint64 {
	return c.user + c.nice + c.system + c.idle + c.iowait + c.irq + c.softirq + c.steal
}

func parseCPU(stat string) (cpuTimes, error) {
	var c cpuTimes
	found := false
	for _, line := range strings.Split(stat, "\n") {
		f := strings.Fields(line)
		if len(f) == 0 || !strings.HasPrefix(f[0], "cpu") {
			continue
		}
		if f[0] != "cpu" {
			c.cpus++
			continue
		}
		if len(f) < 9 {
			return c, fmt.Errorf("/proc/stat: the cpu line has %d fields, not at least 9", len(f))
		}
		v := make([]uint64, 8)
		for i := range v {
			n, err := strconv.ParseUint(f[i+1], 10, 64)
			if err != nil {
				return c, fmt.Errorf("/proc/stat: %q is not a count", f[i+1])
			}
			v[i] = n
		}
		c.user, c.nice, c.system, c.idle, c.iowait, c.irq, c.softirq, c.steal = v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7]
		found = true
	}
	if !found {
		return c, fmt.Errorf("/proc/stat has no cpu line")
	}
	return c, nil
}

// CPU is how all the processors together spent the window, in percent of their time.
type CPU struct {
	CPUs      int     `json:"cpus"`
	IOWaitPct float64 `json:"iowait_pct"`
	UserPct   float64 `json:"user_pct"`
	SystemPct float64 `json:"system_pct"`
	IdlePct   float64 `json:"idle_pct"`
	StealPct  float64 `json:"steal_pct"`
}

func cpuShare(a, b cpuTimes) CPU {
	total := float64(b.total() - a.total())
	pct := func(x, y uint64) float64 {
		if total <= 0 || y < x {
			return 0
		}
		return round1(float64(y-x) / total * 100)
	}
	return CPU{
		CPUs:      b.cpus,
		IOWaitPct: pct(a.iowait, b.iowait),
		UserPct:   pct(a.user+a.nice, b.user+b.nice),
		SystemPct: pct(a.system+a.irq+a.softirq, b.system+b.irq+b.softirq),
		IdlePct:   pct(a.idle, b.idle),
		StealPct:  pct(a.steal, b.steal),
	}
}

// diskCounters is one line of /proc/diskstats (Documentation/admin-guide/iostats.rst). Sectors are
// always 512 bytes there, whatever the device's own sector size.
type diskCounters struct {
	reads, readSectors, readMs    uint64
	writes, writeSectors, writeMs uint64
	inFlight, busyMs, weightedMs  uint64
}

func parseDiskstats(s string) map[string]diskCounters {
	out := map[string]diskCounters{}
	for _, line := range strings.Split(s, "\n") {
		f := strings.Fields(line)
		if len(f) < 14 {
			continue
		}
		n := make([]uint64, 11)
		ok := true
		for i := range n {
			v, err := strconv.ParseUint(f[i+3], 10, 64)
			if err != nil {
				ok = false
				break
			}
			n[i] = v
		}
		if !ok {
			continue
		}
		out[f[2]] = diskCounters{
			reads: n[0], readSectors: n[2], readMs: n[3],
			writes: n[4], writeSectors: n[6], writeMs: n[7],
			inFlight: n[8], busyMs: n[9], weightedMs: n[10],
		}
	}
	return out
}

// Device is one block device over the window.
type Device struct {
	Device      string   `json:"device"`
	Name        string   `json:"name,omitempty"`
	On          []string `json:"on,omitempty"`
	Mounts      []string `json:"mounts,omitempty"`
	BusyPct     float64  `json:"busy_pct"`
	ReadMiBs    float64  `json:"read_mib_s"`
	WriteMiBs   float64  `json:"write_mib_s"`
	ReadIOPS    float64  `json:"read_iops"`
	WriteIOPS   float64  `json:"write_iops"`
	ReadWaitMs  float64  `json:"read_wait_ms"`
	WriteWaitMs float64  `json:"write_wait_ms"`
	QueueDepth  float64  `json:"queue_depth"`
	InFlight    uint64   `json:"in_flight"`
}

func sub(b, a uint64) float64 {
	if b < a {
		return 0
	}
	return float64(b - a)
}

// deviceOver is iostat's arithmetic: busy is the share of the window with any request in flight, wait
// the time a completed request took from queueing to completion, queue depth the requests in flight
// on average.
func deviceOver(name string, a, b diskCounters, window time.Duration) Device {
	ms := float64(window.Milliseconds())
	s := window.Seconds()
	reads, writes := sub(b.reads, a.reads), sub(b.writes, a.writes)
	wait := func(spent, n float64) float64 {
		if n == 0 {
			return 0
		}
		return round1(spent / n)
	}
	busy := sub(b.busyMs, a.busyMs) / ms * 100
	if busy > 100 {
		busy = 100
	}
	return Device{
		Device:      name,
		BusyPct:     round1(busy),
		ReadMiBs:    round1(sub(b.readSectors, a.readSectors) * 512 / mib / s),
		WriteMiBs:   round1(sub(b.writeSectors, a.writeSectors) * 512 / mib / s),
		ReadIOPS:    round1(reads / s),
		WriteIOPS:   round1(writes / s),
		ReadWaitMs:  wait(sub(b.readMs, a.readMs), reads),
		WriteWaitMs: wait(sub(b.writeMs, a.writeMs), writes),
		QueueDepth:  round1(sub(b.weightedMs, a.weightedMs) / ms),
		InFlight:    b.inFlight,
	}
}

// blockDevices are the whole devices in /sys/block, without the ones that are not a disk: loop devices,
// RAM disks and compressed swap in RAM (zram, the memory-pressure module's). all keeps every one, and
// every partition /proc/diskstats names.
func (m *Machine) blockDevices(stats map[string]diskCounters, all bool) []string {
	var out []string
	if all {
		for name := range stats {
			out = append(out, name)
		}
		sort.Strings(out)
		return out
	}
	for _, name := range m.names("/sys/block") {
		if strings.HasPrefix(name, "loop") || strings.HasPrefix(name, "ram") || strings.HasPrefix(name, "zram") {
			continue
		}
		if _, ok := stats[name]; ok {
			out = append(out, name)
		}
	}
	return out
}

// mapperName is a device-mapper device's own name (dm-0 is "root"), "" for any other device.
func (m *Machine) mapperName(dev string) string {
	return strings.TrimSpace(m.read("/sys/block/" + dev + "/dm/name"))
}

// below is what a device-mapper or RAID device is built on.
func (m *Machine) below(dev string) []string { return m.names("/sys/block/" + dev + "/slaves") }

// parentOf is the whole device a partition is on; the device itself when it is whole.
func (m *Machine) parentOf(dev string) string {
	if m.exists("/sys/block/" + dev) {
		return dev
	}
	for _, disk := range m.names("/sys/block") {
		if m.exists("/sys/block/" + disk + "/" + dev) {
			return disk
		}
	}
	return dev
}

// PSI is the kernel's pressure stall information for I/O: the share of time some (or all) runnable
// tasks were stalled waiting on it.
type PSI struct {
	Some PSILine  `json:"some"`
	Full *PSILine `json:"full,omitempty"`
}

// PSILine is one line, its running averages and its share of the sampled window.
type PSILine struct {
	WindowPct float64 `json:"window_pct"`
	Avg10     float64 `json:"avg10_pct"`
	Avg60     float64 `json:"avg60_pct"`
	Avg300    float64 `json:"avg300_pct"`
	totalUs   uint64
}

func parsePSI(s string) (PSI, bool) {
	var p PSI
	some := false
	for _, line := range strings.Split(s, "\n") {
		f := strings.Fields(line)
		if len(f) < 5 {
			continue
		}
		var l PSILine
		for _, kv := range f[1:] {
			k, v, _ := strings.Cut(kv, "=")
			switch k {
			case "avg10":
				l.Avg10, _ = strconv.ParseFloat(v, 64)
			case "avg60":
				l.Avg60, _ = strconv.ParseFloat(v, 64)
			case "avg300":
				l.Avg300, _ = strconv.ParseFloat(v, 64)
			case "total":
				l.totalUs, _ = strconv.ParseUint(v, 10, 64)
			}
		}
		switch f[0] {
		case "some":
			p.Some, some = l, true
		case "full":
			full := l
			p.Full = &full
		}
	}
	return p, some
}

func psiOver(a, b PSI, window time.Duration) PSI {
	share := func(x, y PSILine) PSILine {
		y.WindowPct = round1(sub(y.totalUs, x.totalUs) / float64(window.Microseconds()) * 100)
		if y.WindowPct > 100 {
			y.WindowPct = 100
		}
		return y
	}
	out := PSI{Some: share(a.Some, b.Some)}
	if a.Full != nil && b.Full != nil {
		full := share(*a.Full, *b.Full)
		out.Full = &full
	}
	return out
}