package main

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

// Temperature limits: the owner does not want the machine hotter than a set temperature (CPU 90 °C and GPU 80 °C by
// default). The agent reads both temperatures once a second and, when the recent average is over a limit,
//   - holds new chat requests and jobs back (they wait for the machine to cool, up to a time, then are refused with the reason),
//   - slows the engines it runs itself (ComfyUI, Whisper) by stopping them for part of every second.
// Programs that other services run (Ollama, LM Studio, vLLM) cannot be slowed from here without more rights, so for them
// only the "hold new work" part applies. The settings are a file next to the agent program (like the pause flag), so they
// survive restarts and are changed from the control page (limits.set).

type limitsConfig struct {
	Enabled bool `json:"enabled"`
	CPUMax  int  `json:"cpu_temp_c"` // °C
	GPUMax  int  `json:"gpu_temp_c"` // °C
}

func defaultLimits() limitsConfig { return limitsConfig{Enabled: true, CPUMax: 90, GPUMax: 80} }

const minLimitC, maxLimitC = 50, 100

// limitsFile is a variable for tests.
var limitsFile = func() string {
	exe, err := os.Executable()
	if err != nil {
		return "/tmp/wxch-agent.LIMITS.json"
	}
	return filepath.Join(filepath.Dir(exe), "LIMITS.json")
}

func validLimits(c limitsConfig) error {
	if c.CPUMax < minLimitC || c.CPUMax > maxLimitC || c.GPUMax < minLimitC || c.GPUMax > maxLimitC {
		return fmt.Errorf("the temperature limits must be between %d and %d °C", minLimitC, maxLimitC)
	}
	return nil
}

func loadLimits() limitsConfig {
	c := defaultLimits()
	if b, err := os.ReadFile(limitsFile()); err == nil {
		var x limitsConfig
		if json.Unmarshal(b, &x) == nil && validLimits(x) == nil {
			c = x
		}
	}
	return c
}

func saveLimits(c limitsConfig) error {
	if err := validLimits(c); err != nil {
		return err
	}
	b, _ := json.MarshalIndent(c, "", " ")
	tmp := limitsFile() + ".tmp"
	if err := os.WriteFile(tmp, b, 0644); err != nil {
		return err
	}
	return os.Rename(tmp, limitsFile())
}

// limitCtl is the state of the controller: pure, so that tests can drive it with made-up numbers.
type limitCtl struct {
	cpu, gpu []float64 // the last few samples
	Holding  bool
	Throttle float64 // share of every second the engines are stopped, 0..maxThrottle
	CPUAvg   float64
	GPUAvg   float64
}

// how long a request or a job waits for the load to fall before it is refused
const (
	inferLoadWait = 90 * time.Second
	jobLoadWait   = 10 * time.Minute
)

const (
	limitWindow  = 3   // samples (seconds) averaged
	limitMargin  = 3.0 // the temperature must fall this far below a limit before work is let in again
	maxThrottle  = 0.8
	throttleStep = 0.1
)

func avg(v []float64) float64 {
	if len(v) == 0 {
		return 0
	}
	var s float64
	for _, x := range v {
		s += x
	}
	return s / float64(len(v))
}

// step takes one sample of each temperature and says whether new work is held and how much to slow the engines.
func (l *limitCtl) step(cfg limitsConfig, cpu, gpu float64) {
	l.cpu = append(l.cpu, cpu)
	l.gpu = append(l.gpu, gpu)
	if len(l.cpu) > limitWindow {
		l.cpu, l.gpu = l.cpu[1:], l.gpu[1:]
	}
	l.CPUAvg, l.GPUAvg = avg(l.cpu), avg(l.gpu)
	if !cfg.Enabled {
		l.Holding, l.Throttle = false, 0
		return
	}
	over := l.CPUAvg > float64(cfg.CPUMax) || l.GPUAvg > float64(cfg.GPUMax)
	clear := l.CPUAvg < float64(cfg.CPUMax)-limitMargin && l.GPUAvg < float64(cfg.GPUMax)-limitMargin
	switch {
	case over:
		l.Holding = true
		if l.Throttle += throttleStep; l.Throttle > maxThrottle {
			l.Throttle = maxThrottle
		}
	case clear:
		l.Holding = false
		if l.Throttle -= throttleStep; l.Throttle < 0 {
			l.Throttle = 0
		}
	}
}

var limits = struct {
	sync.Mutex
	cfg limitsConfig
	ctl limitCtl
}{cfg: defaultLimits()}

type limitsStat struct {
	Enabled     bool    `json:"enabled"`
	CPUMax      int     `json:"cpu_temp_max"`
	GPUMax      int     `json:"gpu_temp_max"`
	Holding     bool    `json:"holding"`
	ThrottlePct int     `json:"throttle_pct"`
	CPUAvg      float64 `json:"cpu_temp"` // °C, the recent average
	GPUAvg      float64 `json:"gpu_temp"`
}

func limitsSnapshot() limitsStat {
	limits.Lock()
	defer limits.Unlock()
	return limitsStat{Enabled: limits.cfg.Enabled, CPUMax: limits.cfg.CPUMax, GPUMax: limits.cfg.GPUMax, Holding: limits.ctl.Holding,
		ThrottlePct: int(limits.ctl.Throttle*100 + 0.5), CPUAvg: limits.ctl.CPUAvg, GPUAvg: limits.ctl.GPUAvg}
}

func limitsHolding() (bool, string) {
	s := limitsSnapshot()
	if !s.Holding {
		return false, ""
	}
	return true, fmt.Sprintf("the machine is at its temperature limit (CPU %.0f°C of %d°C, GPU %.0f°C of %d°C)", s.CPUAvg, s.CPUMax, s.GPUAvg, s.GPUMax)
}

// waitForLoad lets new work in as soon as the temperatures are under the limits, and waits for that for at most `max`.
// It returns the reason when the time ran out (or the context ended), "" when the work may start.
func waitForLoad(ctx context.Context, max time.Duration) string {
	deadline := time.Now().Add(max)
	for {
		h, why := limitsHolding()
		if !h {
			return ""
		}
		if time.Now().After(deadline) {
			return why
		}
		select {
		case <-ctx.Done():
			return why
		case <-time.After(500 * time.Millisecond):
		}
	}
}

// gpuTempNow asks the driver for the GPU temperature in °C (0 when it cannot be read).
func gpuTempNow() float64 {
	out, err := runCmd(3*time.Second, "nvidia-smi", "--query-gpu=temperature.gpu", "--format=csv,noheader,nounits")
	if err != nil {
		return 0
	}
	v, _ := strconv.ParseFloat(strings.TrimSpace(strings.SplitN(out, "\n", 2)[0]), 64)
	return v
}

// limitsLoop samples the temperatures every second, updates the controller and slows our own engines when it says so.
func limitsLoop(stop <-chan struct{}) {
	limits.Lock()
	limits.cfg = loadLimits()
	limits.Unlock()
	gpu := 0.0
	for tick := 0; ; tick++ {
		select {
		case <-stop:
			hostPauseAll(false)
			return
		case <-time.After(time.Second):
		}
		cpu := maxThermal() // the hottest thermal zone of the machine (what the dashboard shows as CPU temperature)
		limits.Lock()
		enabled := limits.cfg.Enabled
		limits.Unlock()
		if enabled && tick%2 == 0 { // nvidia-smi is a program start: every other second is enough
			gpu = gpuTempNow()
		}
		limits.Lock()
		was := limits.ctl.Holding
		limits.ctl.step(limits.cfg, cpu, gpu)
		th, holding := limits.ctl.Throttle, limits.ctl.Holding
		limits.Unlock()
		if holding != was {
			logf("temperature limits: %v (CPU %.0f°C, GPU %.0f°C)", map[bool]string{true: "holding new work", false: "released"}[holding], cpu, gpu)
		}
		if th > 0 {
			hostPauseAll(true)
			time.Sleep(time.Duration(th * float64(time.Second)))
			hostPauseAll(false)
		}
	}
}

// limitsRPC implements limits.get / limits.set.
func (r *rpcCtx) limitsRPC() {
	if r.m.Method == "limits.set" {
		var p limitsConfig
		if err := json.Unmarshal(r.m.Params, &p); err != nil {
			r.fail("bad request")
			return
		}
		if err := saveLimits(p); err != nil {
			r.fail("%v", err)
			return
		}
		limits.Lock()
		limits.cfg = p
		if !p.Enabled {
			limits.ctl.Holding, limits.ctl.Throttle = false, 0
		}
		limits.Unlock()
		if !p.Enabled {
			hostPauseAll(false)
		}
		logf("temperature limits set: enabled=%v CPU %d°C GPU %d°C", p.Enabled, p.CPUMax, p.GPUMax)
	}
	r.ok(limitsSnapshot())
}