package metrics

// CPU power, from the RAPL energy counters the kernel keeps under /sys/class/powercap: Intel
// since Sandy Bridge, and AMD Zen through the same driver. Each counter is a running total in
// µJ, so power is the difference between two readings over the time between them. Every
// package is summed (`package-0`, and `package-1` on a two-socket board). The other top-level
// domains are left out. `psys` claims the whole platform but is only as good as the firmware's
// calibration, and on one machine here it reads lower than the package inside it. The MMIO
// copy of a package (`intel-rapl-mmio:0`) measures the same thing again. This is the CPU's own
// figure, shown beside its clock: cores, built-in graphics and memory controller, not disks,
// fans or the power supply.
//
// The counters are root-only since the PLATYPUS side channel (2020). So, like the GPU client
// walk, they are read in-process when the daemon is root, and otherwise by a root helper,
// `hamlab cpu-energy`, started once through sudo or doas and asked over a pipe at every
// sample: one privilege escalation per daemon, not one a second in the auth log. The helper
// takes nothing from the daemon but the newline that asks; it finds and reads the counters
// itself. Without root there is no reading, and nothing is shown.

import (
	"bufio"
	"encoding/json"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"hamlab/internal/priv"
)

const raplRoot = "/sys/class/powercap"

var raplPkgRe = regexp.MustCompile(`^package-\d+$`)

// raplCount is one package's counter at one moment.
type raplCount struct {
	Dir string `json:"d"`
	UJ  uint64 `json:"uj"`
}

type cpuPowerState struct {
	probed    bool
	dirs      []string          // the package domains, found at the first sample
	max       map[string]uint64 // each counter's range: where it wraps to zero
	useHelper bool              // the counters are not readable in-process
	helper    *energyHelper
	prev      []raplCount
	prevAt    time.Time
}

// raplPackages lists the package domains under root, sorted. intel-rapl:0 is a package;
// intel-rapl:0:0 is a part of one (core, uncore, dram).
func raplPackages(root string) []string {
	ents, err := os.ReadDir(root)
	if err != nil {
		return nil
	}
	var out []string
	for _, e := range ents {
		n := e.Name()
		if !strings.HasPrefix(n, "intel-rapl:") || strings.Count(n, ":") != 1 {
			continue
		}
		if raplPkgRe.MatchString(strings.TrimSpace(readStr(filepath.Join(root, n, "name")))) {
			out = append(out, n)
		}
	}
	sort.Strings(out)
	return out
}

func readRAPL(root string, dirs []string) ([]raplCount, bool) {
	out := make([]raplCount, 0, len(dirs))
	for _, d := range dirs {
		b, err := os.ReadFile(filepath.Join(root, d, "energy_uj"))
		if err != nil {
			return nil, false
		}
		v, err := strconv.ParseUint(strings.TrimSpace(string(b)), 10, 64)
		if err != nil {
			return nil, false
		}
		out = append(out, raplCount{d, v})
	}
	return out, len(out) > 0
}

// raplWatts is the power between two readings of the same counters. A counter that wrapped
// is carried across its range. One that went backwards with no known range, or a set of
// counters that changed between readings, gives no answer rather than a wrong one.
func raplWatts(prev, cur []raplCount, max map[string]uint64, dt float64) (float64, bool) {
	if dt <= 0 || len(prev) == 0 || len(prev) != len(cur) {
		return 0, false
	}
	var uj float64
	for i := range cur {
		if cur[i].Dir != prev[i].Dir {
			return 0, false
		}
		if cur[i].UJ >= prev[i].UJ {
			uj += float64(cur[i].UJ - prev[i].UJ)
			continue
		}
		m := max[cur[i].Dir]
		if m == 0 || prev[i].UJ > m {
			return 0, false
		}
		uj += float64(m - prev[i].UJ + cur[i].UJ)
	}
	return uj / 1e6 / dt, true
}

// cpuPower is the CPU packages' power since the previous sample, in watts; 0 when this
// machine has no counters or no way to read them.
func (c *Collector) cpuPower(now time.Time) float64 {
	p := &c.power
	if !p.probed {
		p.probed = true
		p.dirs = raplPackages(raplRoot)
		p.max = map[string]uint64{}
		for _, d := range p.dirs {
			if v, ok := readInt(filepath.Join(raplRoot, d, "max_energy_range_uj")); ok && v > 0 {
				p.max[d] = uint64(v)
			}
		}
	}
	if len(p.dirs) == 0 {
		return 0
	}
	var cur []raplCount
	ok := false
	if !p.useHelper {
		if cur, ok = readRAPL(raplRoot, p.dirs); !ok && priv.Mode() != "root" && priv.Available() {
			p.useHelper = true
		}
	}
	if p.useHelper {
		if p.helper == nil {
			p.helper = &energyHelper{}
		}
		cur, ok = p.helper.read()
	}
	if !ok {
		p.prev = nil
		return 0
	}
	w, good := raplWatts(p.prev, cur, p.max, now.Sub(p.prevAt).Seconds())
	p.prev, p.prevAt = cur, now
	if !good {
		return 0
	}
	return round1(w)
}

// energyHelper is the root child that answers each newline with the package counters, one
// JSON line. The same shape as the GPU client walk's helper, kept separate from it.
type energyHelper struct {
	mu       sync.Mutex
	cmd      *exec.Cmd
	in       io.WriteCloser
	out      *bufio.Reader
	failedAt time.Time
}

func (h *energyHelper) read() ([]raplCount, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.cmd == nil {
		if !h.failedAt.IsZero() && time.Since(h.failedAt) < helperRetry {
			return nil, false
		}
		if !h.start() {
			h.failedAt = time.Now()
			return nil, false
		}
	}
	// a stuck child must not stall the sampler: it is killed, and tried again later
	cmd := h.cmd
	stop := time.AfterFunc(5*time.Second, func() { cmd.Process.Kill() })
	defer stop.Stop()
	if _, err := io.WriteString(h.in, "\n"); err != nil {
		h.fail()
		return nil, false
	}
	line, err := h.out.ReadBytes('\n')
	if err != nil {
		h.fail()
		return nil, false
	}
	var cs []raplCount
	if json.Unmarshal(line, &cs) != nil {
		h.fail()
		return nil, false
	}
	return cs, len(cs) > 0
}

func (h *energyHelper) start() bool {
	self, err := os.Executable()
	if err != nil {
		return false
	}
	full := priv.Wrap(self, "cpu-energy")
	cmd := exec.Command(full[0], full[1:]...)
	in, err := cmd.StdinPipe()
	if err != nil {
		return false
	}
	out, err := cmd.StdoutPipe()
	if err != nil {
		return false
	}
	if err := cmd.Start(); err != nil {
		return false
	}
	h.cmd, h.in, h.out = cmd, in, bufio.NewReader(out)
	return true
}

func (h *energyHelper) fail() {
	if h.cmd != nil {
		h.in.Close()
		h.cmd.Process.Kill()
		go h.cmd.Wait()
	}
	h.cmd, h.in, h.out = nil, nil, nil
	h.failedAt = time.Now()
}

// CPUEnergyHelper is `hamlab cpu-energy`: for every line read on stdin, one JSON line of the
// package counters on stdout. Ends when stdin does — when the daemon that started it has gone.
func CPUEnergyHelper(stdin io.Reader, stdout io.Writer) int {
	return cpuEnergyHelper(raplRoot, stdin, stdout)
}

func cpuEnergyHelper(root string, stdin io.Reader, stdout io.Writer) int {
	dirs := raplPackages(root)
	sc := bufio.NewScanner(stdin)
	w := bufio.NewWriter(stdout)
	for sc.Scan() {
		cs, ok := readRAPL(root, dirs)
		if !ok {
			cs = []raplCount{}
		}
		b, _ := json.Marshal(cs)
		w.Write(b)
		w.WriteByte('\n')
		if w.Flush() != nil {
			return 1
		}
	}
	return 0
}