// mavgpud — the workstation's GPU supervisor.
//
// It runs on the workstation (an AMD 7900 GRE, 16GB), not on homesrv, and it is
// deployed separately from the Maven daemons. Maven does not participate in any
// of this and never asks for a start: it reads /health through internal/llm.Pair
// and either gets the big model or falls back to the resident 1.7B.
//
// The rule, from Vikunja #488: keep llama-server loaded whenever the card is
// free, unload it when it has been idle too long or when another process needs
// the card. Not on demand, because a 7-14B takes tens of seconds to load and a
// world question would be answered by a gap every time the card had been quiet.
// Not always on, because that holds 16GB against the owner's own jobs.
//
// It supervises a second child since 09-08-2026, the CW2 transcriber, and for
// one reason only: it is a ROCm process on the same card. Any GPU service the
// owner leaves running beside this daemon reads as a contender and evicts the
// model, so the card needs one owner rather than two neighbours.
package main

import (
	"context"
	"encoding/json"
	"flag"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"
)

type config struct {
	Listen    string `json:"listen"`     // what Maven talks to
	LlamaAddr string `json:"llama_addr"` // where llama-server binds
	LlamaBin  string `json:"llama_bin"`
	// LlamaArgs must include the flags that bind LlamaAddr. They are passed
	// through untouched so the model, context size and layer count stay the
	// owner's business and not this daemon's schema.
	LlamaArgs []string `json:"llama_args"`

	// Stt is optional. Without it mavgpud supervises llama-server alone, which
	// is everything it did before 09-08-2026.
	Stt *sttConfig `json:"stt,omitempty"`

	KFDRoot   string `json:"kfd_root"`
	DRMDevice string `json:"drm_device"`

	Poll        duration `json:"poll"`
	IdleTimeout duration `json:"idle_timeout"`
	StopGrace   duration `json:"stop_grace"`
	MinFreeVRAM int64    `json:"min_free_vram_bytes"`
	// EvictAfter and StartAfter are counted in polls, not seconds. Both exist
	// to damp flapping: a one-tick blip from a short-lived rocm process must
	// not evict the model, and a card that has just been released must not be
	// grabbed before the previous job has finished unmapping.
	EvictAfter int `json:"evict_after_polls"`
	StartAfter int `json:"start_after_polls"`
}

// sttConfig is the CW2 transcriber, which mavgpud runs for one reason: it is a
// ROCm process on this card. Left to its own systemd unit it registers on the
// KFD, the supervisor reads it as a contender, and llama-server is evicted
// within two polls and restarted five polls later, forever. That thrash was
// observed on 2026-08-09 and it is what folded the service in here.
//
// Maven talks to it directly, not through this daemon. There is no proxy and no
// idle timer: at 1.6GB it denies the card to nobody, and unloading it would only
// send the next voice turn to the homesrv floor for no gain.
type sttConfig struct {
	// Addr is where the service binds, and it is read only to probe /health.
	Addr string   `json:"addr"`
	Bin  string   `json:"bin"`
	Args []string `json:"args"`
}

func defaults() config {
	return config{
		Listen:      ":8080",
		LlamaAddr:   "127.0.0.1:8081",
		KFDRoot:     "/sys/class/kfd/kfd/proc",
		DRMDevice:   "/sys/class/drm/card1/device",
		Poll:        duration(time.Second),
		IdleTimeout: duration(15 * time.Minute),
		StopGrace:   duration(20 * time.Second),
		MinFreeVRAM: 15 << 30,
		EvictAfter:  2,
		StartAfter:  5,
	}
}

// duration lets the config file say "15m" instead of counting nanoseconds.
type duration time.Duration

func (d *duration) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	v, err := time.ParseDuration(s)
	if err != nil {
		return err
	}
	*d = duration(v)
	return nil
}

func main() {
	path := flag.String("config", "/etc/mavgpud.json", "config file")
	flag.Parse()

	cfg := defaults()
	b, err := os.ReadFile(*path)
	if err != nil {
		log.Fatalf("mavgpud: read config: %v", err)
	}
	if err := json.Unmarshal(b, &cfg); err != nil {
		log.Fatalf("mavgpud: parse config: %v", err)
	}
	if cfg.LlamaBin == "" {
		log.Fatal("mavgpud: llama_bin is required")
	}

	base := "http://" + cfg.LlamaAddr
	run := newRunner("llama-server", cfg.LlamaBin, cfg.LlamaArgs, base+"/health")
	sup := &supervisor{
		cfg:   cfg,
		probe: probe{kfdRoot: cfg.KFDRoot, drmDev: cfg.DRMDevice},
		run:   run,
	}
	if s := cfg.Stt; s != nil {
		if s.Bin == "" || s.Addr == "" {
			log.Fatal("mavgpud: stt needs both bin and addr")
		}
		sup.stt = newRunner("cw2", s.Bin, s.Args, "http://"+s.Addr+"/health")
	}
	sup.touch()

	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer cancel()

	target, err := url.Parse(base)
	if err != nil {
		log.Fatalf("mavgpud: llama_addr: %v", err)
	}
	srv := &http.Server{Addr: cfg.Listen, Handler: sup.handler(target)}
	go func() {
		log.Printf("mavgpud: listening on %s, model %s", cfg.Listen, cfg.LlamaBin)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("mavgpud: listen: %v", err)
		}
	}()

	sup.loop(ctx)

	// The card must come back before we do. A supervisor that exits leaving
	// llama-server holding 14GB is worse than one that never ran.
	shut, done := context.WithTimeout(context.Background(), 5*time.Second)
	defer done()
	_ = srv.Shutdown(shut)
	for _, r := range sup.children() {
		r.stop(time.Duration(cfg.StopGrace))
	}
}

type supervisor struct {
	cfg   config
	probe probe
	run   *runner
	// stt is the CW2 transcriber, or nil when the config names none.
	stt *runner

	lastReq atomic.Int64 // unix nanos of the last request Maven sent

	foreignStreak int
	clearStreak   int
}

func (s *supervisor) touch() { s.lastReq.Store(time.Now().UnixNano()) }

func (s *supervisor) idle() time.Duration {
	return time.Since(time.Unix(0, s.lastReq.Load()))
}

// handler serves the two things the workstation exposes.
//
// /health is answered locally and always, with no GPU cost and no round trip,
// because it is the only thing Maven reads and Maven reads it on a timer
// forever. Everything else is llama-server's API, reverse-proxied. Proxying
// rather than pointing Maven straight at llama-server is what makes the idle
// window measurable: the supervisor cannot otherwise know when the model was
// last used.
func (s *supervisor) handler(target *url.URL) http.Handler {
	proxy := httputil.NewSingleHostReverseProxy(target)
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		if !s.run.isReady() {
			http.Error(w, "model not loaded", http.StatusServiceUnavailable)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"status":"ok"}`))
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if !s.run.isReady() {
			http.Error(w, "model not loaded", http.StatusServiceUnavailable)
			return
		}
		s.touch()
		proxy.ServeHTTP(w, r)
	})
	return mux
}

func (s *supervisor) loop(ctx context.Context) {
	t := time.NewTicker(time.Duration(s.cfg.Poll))
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			s.tick(ctx)
		}
	}
}

// tick is the whole decision. Yielding is checked before starting, and presence
// on the KFD is what triggers it — not a VRAM threshold. A ROCm process
// registers under /sys/class/kfd/kfd/proc when it initialises HIP, before it
// allocates, so we see a contender during its startup rather than after it has
// already failed to get the memory it wanted.
func (s *supervisor) tick(ctx context.Context) {
	var pids []int
	for _, r := range s.children() {
		pids = append(pids, r.pid())
	}
	others := s.probe.foreign(pids...)
	if len(others) > 0 {
		s.foreignStreak++
		s.clearStreak = 0
	} else {
		s.foreignStreak = 0
		s.clearStreak++
	}

	// Yielding is all or nothing. A CPT run wants the whole card, and handing
	// back 8GB while holding 1.6GB is the shape of a failed allocation.
	if s.foreignStreak >= s.cfg.EvictAfter && s.anyRunning() {
		log.Printf("mavgpud: yielding the card to %s", describe(others))
		for _, r := range s.children() {
			r.stop(time.Duration(s.cfg.StopGrace))
		}
		return
	}

	clear := s.clearStreak >= s.cfg.StartAfter

	if s.run.running() {
		s.run.refreshReady(ctx)
		if s.idle() > time.Duration(s.cfg.IdleTimeout) {
			log.Printf("mavgpud: idle for %s, unloading", s.idle().Round(time.Second))
			s.run.stop(time.Duration(s.cfg.StopGrace))
		}
	} else if clear && s.probe.freeVRAM() >= s.cfg.MinFreeVRAM {
		s.touch() // the idle clock starts at load, not at the last request before it
		if err := s.run.start(); err != nil {
			log.Printf("mavgpud: start llama-server: %v", err)
		}
	}

	if s.stt == nil {
		return
	}
	if s.stt.running() {
		s.stt.refreshReady(ctx)
		return
	}
	// No VRAM precondition here, unlike llama-server. That check exists because
	// a 12B refuses to load when the card is short, and 1.6GB fits wherever the
	// KFD is clear. Reading free VRAM would also block the transcriber for good
	// once the language model was resident, since it holds more than the floor.
	if clear {
		if err := s.stt.start(); err != nil {
			log.Printf("mavgpud: start cw2: %v", err)
		}
	}
}

func (s *supervisor) children() []*runner {
	if s.stt == nil {
		return []*runner{s.run}
	}
	return []*runner{s.run, s.stt}
}

func (s *supervisor) anyRunning() bool {
	for _, r := range s.children() {
		if r.running() {
			return true
		}
	}
	return false
}

// describe names the contenders in the log. This log is the instrument for the
// open question in #488: whether polling the KFD misses a job that wants the
// card without registering there.
func describe(procs []gpuProc) string {
	out := ""
	for i, p := range procs {
		if i > 0 {
			out += ", "
		}
		out += p.Comm
	}
	return out
}