package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// The power the machine draws, measured honestly and said with what was measured (novox/hq ADR 0255).
//
//   - **On battery,** the battery's discharge rate. That is the whole machine.
//   - **Otherwise** — on mains, or on a machine without a battery — the processor package's power from
//     the kernel's powercap energy counter (RAPL, on Intel and AMD alike: a rate from the counter's
//     change), plus a discrete graphics device's own power where it reports one. That is the processor
//     (and graphics) part, never called the machine's.
//   - **The machine at the wall** needs a metered plug or a UPS, and is not measured here.
//
// The energy counter is root's to read, on purpose: read often, it leaks what the processor is doing
// (the PLATYPUS attack). So the counter is never opened to the operator's account. `power draw` runs as
// root, as the module's process, samples every SampleEvery, and writes only the averaged rate to
// DrawFile, too coarse to leak anything. That file is this module's alone: its bundle reads it and puts
// the reading on the bus as the state `draw`, keyed by the machine (publish.go), and nothing of another
// module reads it (novox/hq ADR 0260: modules talk over the bus).

// DrawFile is where the sampler leaves the reading for the module's bundle, and SampleEvery how often.
const (
	DrawFile    = "/run/mesh-power/draw"
	SampleEvery = 5 * time.Second
)

// Draw is one reading: each part measured, in watts, and what it covers.
type Draw struct {
	// Machine is the battery's discharge, when on battery: the whole machine.
	Machine *float64
	// CPU is the processor packages' power, from the energy counters.
	CPU *float64
	// GPU is the discrete graphics devices' own power, where they report it; asleep counts as none.
	GPU *float64
}

// Reading is a draw as the state `draw` holds it: the line shown, the watts, what they cover and when.
type Reading struct {
	Text     string  `json:"text"`
	Watts    float64 `json:"watts,omitempty"`
	Measures string  `json:"measures,omitempty"`
	At       string  `json:"at"`
}

// Reading is the draw as it is published, at a moment.
func (d Draw) Reading(at time.Time) Reading {
	r := Reading{Text: d.Line(), At: at.UTC().Format(time.RFC3339)}
	switch {
	case d.Machine != nil:
		r.Watts, r.Measures = *d.Machine, "machine"
	case d.CPU != nil && d.GPU != nil:
		r.Watts, r.Measures = *d.CPU+*d.GPU, "cpu+gpu"
	case d.CPU != nil:
		r.Watts, r.Measures = *d.CPU, "cpu"
	case d.GPU != nil:
		r.Watts, r.Measures = *d.GPU, "gpu"
	}
	return r
}

// Line is the reading as the bar shows it.
func (d Draw) Line() string {
	switch {
	case d.Machine != nil:
		return fmt.Sprintf("%.1f W machine", *d.Machine)
	case d.CPU != nil && d.GPU != nil:
		return fmt.Sprintf("%.1f W CPU+GPU", *d.CPU+*d.GPU)
	case d.CPU != nil:
		return fmt.Sprintf("%.1f W CPU", *d.CPU)
	case d.GPU != nil:
		return fmt.Sprintf("%.1f W GPU", *d.GPU)
	}
	return "power not measured"
}

// counter is one energy counter's reading, in microjoules, with the range it wraps at.
type counter struct {
	uj, max int64
}

// MaxPackageWatts is more than any processor package draws. A rate above it is a counter that was
// reset — after a resume, say — not power, and that sample gives no reading.
const MaxPackageWatts = 1000

// packageCounters is every processor package's energy counter, by the package's name, each counted
// once. Intel exposes some packages twice, through MSRs (`intel-rapl:N`) and through MMIO
// (`intel-rapl-mmio:N`), under one name; the MSR one is read, and the second never adds to it.
func (m *Machine) packageCounters() map[string]counter {
	out := map[string]counter{}
	from := map[string]string{}
	for _, dir := range m.glob("/sys/class/powercap/*") {
		name := m.read(dir + "/name")
		if !strings.HasPrefix(name, "package") {
			continue
		}
		if was, seen := from[name]; seen && !strings.Contains(filepath.Base(was), "mmio") {
			continue
		}
		uj, ok := m.readInt(dir + "/energy_uj")
		if !ok {
			continue
		}
		max, _ := m.readInt(dir + "/max_energy_range_uj")
		out[name] = counter{uj: uj, max: max}
		from[name] = dir
	}
	return out
}

// cpuWatts is the packages' power between two readings `elapsed` apart; nil when no counter was read
// both times. A counter that wrapped is counted across its range; one that gives a rate no package
// draws was reset, and the sample gives nothing rather than a number that is not power.
func cpuWatts(before, after map[string]counter, elapsed time.Duration) *float64 {
	if elapsed <= 0 {
		return nil
	}
	var joules float64
	read := false
	for dir, a := range after {
		b, ok := before[dir]
		if !ok {
			continue
		}
		delta := a.uj - b.uj
		if delta < 0 {
			if a.max <= 0 {
				continue
			}
			delta += a.max
		}
		if float64(delta)/1e6/elapsed.Seconds() > MaxPackageWatts {
			return nil
		}
		joules += float64(delta) / 1e6
		read = true
	}
	if !read {
		return nil
	}
	w := joules / elapsed.Seconds()
	return &w
}

// batteryWatts is the system battery's discharge, when it is discharging; nil otherwise.
func (m *Machine) batteryWatts() *float64 {
	var total float64
	found := false
	for _, s := range m.Supplies() {
		if s.Type != "Battery" || s.Device || s.Status != "Discharging" {
			continue
		}
		dir := "/sys/class/power_supply/" + s.Name
		if p, ok := m.readInt(dir + "/power_now"); ok {
			total += float64(abs(p)) / 1e6
			found = true
			continue
		}
		i, iok := m.readInt(dir + "/current_now")
		v, vok := m.readInt(dir + "/voltage_now")
		if iok && vok {
			total += float64(abs(i)) * float64(v) / 1e12
			found = true
		}
	}
	if !found {
		return nil
	}
	return &total
}

func abs(n int64) int64 {
	if n < 0 {
		return -n
	}
	return n
}

// gpuWatts is the discrete graphics devices' own power, where they report it; nil when none does.
//
//   - An AMD device's hwmon power counts only when it has its own memory (`mem_info_vram_vendor`): on
//     an APU the integrated graphics report the whole package, which the processor's counter already
//     holds, and counting it twice would be a lie.
//   - An NVIDIA device is asked through nvidia-smi only while it is awake: asking wakes a suspended GPU,
//     which on a hybrid laptop costs more than the answer is worth. Asleep, it draws next to nothing and
//     counts as none.
func (m *Machine) gpuWatts(ctx context.Context) *float64 {
	var total float64
	found := false
	for _, card := range m.glob("/sys/class/drm/card[0-9]*") {
		if strings.Contains(filepath.Base(card), "-") || !m.exists(card+"/device/mem_info_vram_vendor") {
			continue
		}
		// Asleep, it is not read: reading its sensors wakes it, as asking nvidia-smi would.
		if status := m.read(card + "/device/power/runtime_status"); status != "" && status != "active" {
			continue
		}
		for _, hw := range m.glob(card + "/device/hwmon/hwmon*") {
			uw, ok := m.readInt(hw + "/power1_average")
			if !ok {
				uw, ok = m.readInt(hw + "/power1_input")
			}
			if ok {
				total += float64(uw) / 1e6
				found = true
				break
			}
		}
	}
	for _, gpu := range m.glob("/proc/driver/nvidia/gpus/*") {
		address := filepath.Base(gpu)
		if m.read("/sys/bus/pci/devices/"+address+"/power/runtime_status") != "active" {
			continue
		}
		out, err := m.Run(ctx, "nvidia-smi", "--query-gpu=power.draw", "--format=csv,noheader,nounits", "-i", address)
		if err != nil {
			continue
		}
		if w, err := strconv.ParseFloat(strings.TrimSpace(out), 64); err == nil {
			total += w
			found = true
		}
	}
	if !found {
		return nil
	}
	return &total
}

// Sample is one reading over `every`: the battery's discharge if the machine is on battery, else the
// processor's and the graphics' parts.
func (m *Machine) Sample(ctx context.Context, every time.Duration, wait func(time.Duration) bool) Draw {
	if w := m.batteryWatts(); w != nil {
		return Draw{Machine: w}
	}
	before, start := m.packageCounters(), time.Now()
	if !wait(every) {
		return Draw{}
	}
	return Draw{CPU: cpuWatts(before, m.packageCounters(), time.Since(start)), GPU: m.gpuWatts(ctx)}
}

// Publish writes the reading where the module's bundle reads it, whole, replacing the last.
func (m *Machine) Publish(r Reading) error {
	raw, err := json.Marshal(r)
	if err != nil {
		return err
	}
	line := string(raw)
	path := m.path(DrawFile)
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	tmp := path + ".new"
	if err := os.WriteFile(tmp, []byte(line+"\n"), 0o644); err != nil {
		return err
	}
	if err := os.Chmod(tmp, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

// RunDraw samples and publishes until the context ends.
func RunDraw(ctx context.Context, m *Machine) error {
	return runDraw(ctx, m, SampleEvery, m.Publish)
}

// runDraw publishes one reading every `every`. A battery reading is instant, so it waits the interval
// itself; a counter reading waits inside its sample. It returns as soon as the context ends.
func runDraw(ctx context.Context, m *Machine, every time.Duration, publish func(Reading) error) error {
	wait := func(d time.Duration) bool {
		select {
		case <-ctx.Done():
			return false
		case <-time.After(d):
			return true
		}
	}
	for ctx.Err() == nil {
		d := m.Sample(ctx, every, wait)
		if ctx.Err() != nil {
			break
		}
		if err := publish(d.Reading(time.Now())); err != nil {
			return fmt.Errorf("publishing the draw: %w", err)
		}
		if d.Machine != nil && !wait(every) {
			break
		}
	}
	return nil
}