package main

import (
	"context"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
)

// procIO is what /proc/<pid>/io says a process caused to be read from and written to storage: not
// what it read from the page cache, which rchar and wchar count.
type procIO struct {
	read, write uint64
}

// processIO reads every process's storage counters. A process of another account is unreadable unless
// this runtime runs as root; those are counted, never guessed.
func (m *Machine) processIO() (map[int]procIO, int) {
	out := map[int]procIO{}
	unreadable := 0
	for _, name := range m.names("/proc") {
		pid, err := strconv.Atoi(name)
		if err != nil {
			continue
		}
		text := m.read("/proc/" + name + "/io")
		if text == "" {
			if m.exists("/proc/" + name) {
				unreadable++
			}
			continue
		}
		var p procIO
		for _, line := range strings.Split(text, "\n") {
			k, v, ok := strings.Cut(line, ":")
			if !ok {
				continue
			}
			n, _ := strconv.ParseUint(strings.TrimSpace(v), 10, 64)
			switch k {
			case "read_bytes":
				p.read = n
			case "write_bytes":
				p.write = n
			}
		}
		out[pid] = p
	}
	return out, unreadable
}

// Process is one process's storage I/O over the window, or a unit's or a container's (with its requests
// per second, which io.stat counts and /proc/<pid>/io does not).
type Process struct {
	PID       int     `json:"pid,omitempty"`
	Command   string  `json:"command,omitempty"`
	Unit      string  `json:"unit,omitempty"`
	Container string  `json:"container,omitempty"`
	ReadMiBs  float64 `json:"read_mib_s"`
	WriteMiBs float64 `json:"write_mib_s"`
	ReadMiB   float64 `json:"read_mib"`
	WriteMiB  float64 `json:"write_mib"`
	ReadIOPS  float64 `json:"read_iops,omitempty"`
	WriteIOPS float64 `json:"write_iops,omitempty"`
	bytes     float64
}

var containerID = regexp.MustCompile(`(?:docker-|/docker/)([0-9a-f]{64})`)

// placeOf is the systemd unit and the container a process runs in, from its cgroup.
func (m *Machine) placeOf(pid int) (unit, container string) {
	for _, line := range strings.Split(m.read("/proc/"+strconv.Itoa(pid)+"/cgroup"), "\n") {
		parts := strings.SplitN(line, ":", 3)
		if len(parts) != 3 || (parts[0] != "0" && !strings.Contains(parts[1], "name=systemd")) {
			continue
		}
		path := strings.TrimSpace(parts[2])
		if id := containerID.FindStringSubmatch(path); id != nil {
			container = id[1]
		}
		segs := strings.Split(path, "/")
		for i := len(segs) - 1; i >= 0; i-- {
			if strings.HasSuffix(segs[i], ".service") || strings.HasSuffix(segs[i], ".scope") {
				unit = segs[i]
				break
			}
		}
		if parts[0] == "0" {
			break
		}
	}
	return unit, container
}

// containerNames is each running container's full id with its name, from the container runtime; none
// when it cannot be asked, and then a container is said by its short id.
func (m *Machine) containerNames() map[string]string {
	out := map[string]string{}
	if m.Run == nil {
		return out
	}
	text, err := m.Run(context.Background(), "docker", "ps", "--no-trunc", "--format", "{{.ID}}\t{{.Names}}")
	if err != nil {
		return out
	}
	for _, line := range strings.Split(text, "\n") {
		id, name, ok := strings.Cut(strings.TrimSpace(line), "\t")
		if ok {
			out[id] = name
		}
	}
	return out
}

// topProcesses is the processes that moved the most bytes to and from storage between the two
// readings, the largest first, each with the unit and the container it runs in.
func (m *Machine) topProcesses(a, b map[int]procIO, window time.Duration, limit int) []Process {
	var moved []Process
	for pid, end := range b {
		start, ok := a[pid]
		if !ok {
			continue
		}
		r, w := sub(end.read, start.read), sub(end.write, start.write)
		if r+w == 0 {
			continue
		}
		// Bytes until inMiB says them in MiB.
		moved = append(moved, Process{PID: pid, ReadMiB: r, WriteMiB: w, bytes: r + w})
	}
	sort.Slice(moved, func(i, j int) bool {
		if moved[i].bytes != moved[j].bytes {
			return moved[i].bytes > moved[j].bytes
		}
		return moved[i].PID < moved[j].PID
	})
	if len(moved) > limit {
		moved = moved[:limit]
	}
	var names map[string]string
	for i := range moved {
		p := &moved[i]
		p.Command = strings.TrimSpace(m.read("/proc/" + strconv.Itoa(p.PID) + "/comm"))
		var id string
		p.Unit, id = m.placeOf(p.PID)
		if id != "" {
			if names == nil {
				names = m.containerNames()
			}
			p.Container = id[:12]
			if n, ok := names[id]; ok {
				p.Container = n
			}
		}
		inMiB(p, window.Seconds())
	}
	if moved == nil {
		moved = []Process{}
	}
	return moved
}

// inMiB turns a reading's bytes into MiB and MiB/s, to two decimals: a process writing a log moves
// kilobytes, and one decimal of a MiB would say it moved nothing.
func inMiB(p *Process, s float64) {
	r, w := p.ReadMiB, p.WriteMiB
	p.ReadMiBs, p.WriteMiBs = round2(r/mib/s), round2(w/mib/s)
	p.ReadMiB, p.WriteMiB = round2(r/mib), round2(w/mib)
}

func round2(f float64) float64 { return float64(int64(f*100+0.5)) / 100 }