package workerload

import (
	"errors"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

type SystemSampler struct {
	mu sync.Mutex

	root          string
	lastCPUUsage  int64
	lastCPUTime   time.Time
	lastThrottled int64
}

func NewSystemSampler() *SystemSampler {
	return &SystemSampler{root: "/sys/fs/cgroup"}
}

func NewSystemSamplerAt(root string) *SystemSampler {
	return &SystemSampler{root: strings.TrimRight(root, "/")}
}

func (s *SystemSampler) Sample(databaseConnections, databaseMax int32) ResourceSample {
	s.mu.Lock()
	defer s.mu.Unlock()
	now := time.Now()
	memoryCurrent, memoryLimit := s.memory()
	cpuUsage, throttledCount := s.cpuStat()
	cpuLimit := s.cpuLimit()
	cpuUtilization := 0.0
	if s.lastCPUUsage > 0 && cpuUsage >= s.lastCPUUsage && !s.lastCPUTime.IsZero() {
		elapsed := now.Sub(s.lastCPUTime).Seconds()
		if elapsed > 0 {
			cpuUtilization = float64(cpuUsage-s.lastCPUUsage) / 1_000_000 / elapsed / cpuLimit
		}
	}
	throttled := s.lastThrottled > 0 && throttledCount > s.lastThrottled
	s.lastCPUUsage = cpuUsage
	s.lastThrottled = throttledCount
	s.lastCPUTime = now
	return ResourceSample{
		MemoryCurrentBytes: memoryCurrent,
		MemoryLimitBytes:   memoryLimit,
		CPUUtilization:     bounded(cpuUtilization),
		CPUThrottled:       throttled,
		DBConnections:      databaseConnections,
		DBMaxConnections:   databaseMax,
		SampledAt:          now,
	}
}

func (s *SystemSampler) memory() (int64, int64) {
	current, currentErr := readIntFile(s.root + "/memory.current")
	limit, limitErr := readLimitFile(s.root + "/memory.max")
	if currentErr == nil && limitErr == nil {
		return current, limit
	}
	current, _ = readIntFile(s.root + "/memory/memory.usage_in_bytes")
	limit, _ = readLimitFile(s.root + "/memory/memory.limit_in_bytes")
	return current, limit
}

func (s *SystemSampler) cpuStat() (usageUsec, throttled int64) {
	data, err := os.ReadFile(s.root + "/cpu.stat")
	if err != nil {
		return 0, 0
	}
	for _, line := range strings.Split(string(data), "\n") {
		fields := strings.Fields(line)
		if len(fields) != 2 {
			continue
		}
		value, parseErr := strconv.ParseInt(fields[1], 10, 64)
		if parseErr != nil {
			continue
		}
		switch fields[0] {
		case "usage_usec":
			usageUsec = value
		case "nr_throttled":
			throttled = value
		}
	}
	return usageUsec, throttled
}

func (s *SystemSampler) cpuLimit() float64 {
	data, err := os.ReadFile(s.root + "/cpu.max")
	if err == nil {
		fields := strings.Fields(string(data))
		if len(fields) == 2 && fields[0] != "max" {
			quota, quotaErr := strconv.ParseFloat(fields[0], 64)
			period, periodErr := strconv.ParseFloat(fields[1], 64)
			if quotaErr == nil && periodErr == nil && quota > 0 && period > 0 {
				return max(quota/period, .001)
			}
		}
	}
	return max(float64(runtime.GOMAXPROCS(0)), 1)
}

func readIntFile(path string) (int64, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}
	return strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64)
}

func readLimitFile(path string) (int64, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}
	value := strings.TrimSpace(string(data))
	if value == "" || value == "max" {
		return 0, errors.New("cgroup limit is unlimited")
	}
	limit, err := strconv.ParseInt(value, 10, 64)
	if err != nil || limit <= 0 || limit > 1<<60 {
		return 0, errors.New("cgroup limit is not finite")
	}
	return limit, nil
}