package discovery

import (
	"bufio"
	"bytes"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Load is how busy a node is: the machine as a whole, and each engine's share.
type Load struct {
	CPUPercent float64
	CPUCount   int
	Load1      float64
	Load5      float64
	Load15     float64
	MemUsed    int64
	MemAvail   int64
	Shards     []ShardLoad
}

// ShardLoad is one engine process group's use.
type ShardLoad struct {
	ID       string
	CPUCores float64
	RSSBytes int64
}

// userHZ is the unit /proc counts CPU time in. Fixed at 100 by the kernel's
// ABI on every architecture a node runs, whatever the internal tick rate.
const userHZ = 100

// LoadSampler measures load as the difference between two readings of /proc.
//
// CPU time is a counter, so a percentage needs two samples and the time
// between them. The sampler keeps the previous one; the first call has
// nothing to compare with and reports no load at all, rather than a zero that
// would read as an idle machine.
//
// Linux only, like every node. Elsewhere the files are missing and it reports
// nothing, as an older agent does.
type LoadSampler struct {
	root string
	now  func() time.Time

	mu      sync.Mutex
	prevAt  time.Time
	prevCPU cpuTimes
	prevPG  map[int]groupUse // process group -> what it has used
}

// groupUse is a process group's CPU time so far and its resident memory now.
type groupUse struct {
	jiffies  uint64
	rssPages int64
}

type cpuTimes struct{ total, idle uint64 }

// NewLoadSampler reads the real /proc.
func NewLoadSampler() *LoadSampler { return newLoadSampler("/proc", time.Now) }

func newLoadSampler(root string, now func() time.Time) *LoadSampler {
	return &LoadSampler{root: root, now: now}
}

// Sample reads the machine and the process groups led by the given shards'
// PIDs, keyed by shard ID. False when there is nothing to report yet.
func (s *LoadSampler) Sample(shards map[string]int) (Load, bool) {
	cpu, ok := s.cpu()
	if !ok {
		return Load{}, false
	}
	groups := s.groups()
	at := s.now()

	s.mu.Lock()
	prevAt, prevCPU, prevPG := s.prevAt, s.prevCPU, s.prevPG
	s.prevAt, s.prevCPU, s.prevPG = at, cpu, groups
	s.mu.Unlock()

	if prevAt.IsZero() || cpu.total <= prevCPU.total {
		return Load{}, false
	}
	l := Load{CPUCount: s.cpuCount()}
	dTotal := float64(cpu.total - prevCPU.total)
	dIdle := float64(cpu.idle - prevCPU.idle)
	l.CPUPercent = clamp((dTotal-dIdle)/dTotal*100, 0, 100)
	l.Load1, l.Load5, l.Load15 = s.loadavg()
	total, avail := s.meminfo()
	if total > 0 {
		l.MemUsed, l.MemAvail = total-avail, avail
	}

	secs := at.Sub(prevAt).Seconds()
	for id, pid := range shards {
		if pid <= 0 {
			continue
		}
		now := groups[pid]
		sl := ShardLoad{ID: id, RSSBytes: now.rssPages * int64(os.Getpagesize())}
		if before, ok := prevPG[pid]; ok && secs > 0 && now.jiffies >= before.jiffies {
			sl.CPUCores = float64(now.jiffies-before.jiffies) / userHZ / secs
		}
		l.Shards = append(l.Shards, sl)
	}
	return l, true
}

func clamp(v, lo, hi float64) float64 { return max(lo, min(v, hi)) }

// cpu reads the aggregate line of /proc/stat. Idle includes iowait: a core
// waiting on the disk is not doing work.
func (s *LoadSampler) cpu() (cpuTimes, bool) {
	b, err := os.ReadFile(filepath.Join(s.root, "stat"))
	if err != nil {
		return cpuTimes{}, false
	}
	line, _, _ := bytes.Cut(b, []byte("\n"))
	f := strings.Fields(string(line))
	if len(f) < 5 || f[0] != "cpu" {
		return cpuTimes{}, false
	}
	var t cpuTimes
	for i, v := range f[1:] {
		// guest and guest_nice are already counted in user and nice.
		if i >= 8 {
			break
		}
		n, _ := strconv.ParseUint(v, 10, 64)
		t.total += n
		if i == 3 || i == 4 {
			t.idle += n
		}
	}
	return t, true
}

func (s *LoadSampler) cpuCount() int {
	b, err := os.ReadFile(filepath.Join(s.root, "stat"))
	if err != nil {
		return 0
	}
	n := 0
	for _, line := range strings.Split(string(b), "\n") {
		if strings.HasPrefix(line, "cpu") && len(line) > 3 && line[3] >= '0' && line[3] <= '9' {
			n++
		}
	}
	return n
}

func (s *LoadSampler) loadavg() (l1, l5, l15 float64) {
	b, err := os.ReadFile(filepath.Join(s.root, "loadavg"))
	if err != nil {
		return 0, 0, 0
	}
	f := strings.Fields(string(b))
	if len(f) < 3 {
		return 0, 0, 0
	}
	l1, _ = strconv.ParseFloat(f[0], 64)
	l5, _ = strconv.ParseFloat(f[1], 64)
	l15, _ = strconv.ParseFloat(f[2], 64)
	return l1, l5, l15
}

// meminfo returns MemTotal and MemAvailable in bytes.
func (s *LoadSampler) meminfo() (total, avail int64) {
	f, err := os.Open(filepath.Join(s.root, "meminfo"))
	if err != nil {
		return 0, 0
	}
	defer f.Close()
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		k, v, ok := strings.Cut(sc.Text(), ":")
		if !ok {
			continue
		}
		fields := strings.Fields(v)
		if len(fields) == 0 {
			continue
		}
		kb, _ := strconv.ParseInt(fields[0], 10, 64)
		switch k {
		case "MemTotal":
			total = kb * 1024
		case "MemAvailable":
			avail = kb * 1024
		}
	}
	return total, avail
}

// procStat is what one /proc/<pid>/stat says that load needs.
type procStat struct {
	pgrp     int
	jiffies  uint64
	rssPages int64
}

// readProcStat parses /proc/<pid>/stat. The command name is in parentheses
// and may itself hold spaces and parentheses, so fields are counted from the
// last closing one.
func readProcStat(path string) (procStat, bool) {
	b, err := os.ReadFile(path)
	if err != nil {
		return procStat{}, false
	}
	i := bytes.LastIndexByte(b, ')')
	if i < 0 {
		return procStat{}, false
	}
	f := strings.Fields(string(b[i+1:]))
	// After the name: state(0) ppid(1) pgrp(2) ... utime(11) stime(12) ... rss(21)
	if len(f) < 22 {
		return procStat{}, false
	}
	var st procStat
	st.pgrp, _ = strconv.Atoi(f[2])
	ut, _ := strconv.ParseUint(f[11], 10, 64)
	stt, _ := strconv.ParseUint(f[12], 10, 64)
	st.jiffies = ut + stt
	st.rssPages, _ = strconv.ParseInt(f[21], 10, 64)
	return st, true
}

// groups totals CPU time and memory per process group, in one pass over
// /proc. An engine is started as the leader of its own group, so its workers
// — vLLM starts several — are counted with it.
func (s *LoadSampler) groups() map[int]groupUse {
	out := map[int]groupUse{}
	s.eachProc(func(st procStat) {
		g := out[st.pgrp]
		g.jiffies += st.jiffies
		g.rssPages += st.rssPages
		out[st.pgrp] = g
	})
	return out
}

func (s *LoadSampler) eachProc(fn func(procStat)) {
	entries, err := os.ReadDir(s.root)
	if err != nil {
		return
	}
	for _, e := range entries {
		if _, err := strconv.Atoi(e.Name()); err != nil {
			continue
		}
		if st, ok := readProcStat(filepath.Join(s.root, e.Name(), "stat")); ok {
			fn(st)
		}
	}
}