// gpu-turnstile is a GPU arbitration proxy that sits in front of Ollama and
// ComfyUI and guarantees only one of them uses the GPU at a time.
package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"gpu-turnstile/internal/comfy"
	"gpu-turnstile/internal/config"
	"gpu-turnstile/internal/game"
	"gpu-turnstile/internal/lock"
	"gpu-turnstile/internal/metrics"
	"gpu-turnstile/internal/ollama"
	"gpu-turnstile/internal/proxy"
	"gpu-turnstile/internal/service"
	"gpu-turnstile/internal/supervise"
	"gpu-turnstile/internal/update"
)

// version is injected at build time via -ldflags "-X main.version=...".
var version = "dev"

// exitCodeUpdate tells the service recovery configuration to restart the
// process: a signed update has been staged and the GPU lock is idle.
const exitCodeUpdate = 3

// stdoutIsTerminal reports whether stdout is a console (char device), as
// opposed to a pipe or file — which is what Docker containers and services
// see.
func stdoutIsTerminal() bool {
	fi, err := os.Stdout.Stat()
	return err == nil && fi.Mode()&os.ModeCharDevice != 0
}

// parseFlags extracts -config <path> (or -config=<path>), the
// --install-service / --remove-service switches, --no-copy, -h/--help,
// -v/--version, --force-update and the hidden --elevated-child marker from
// args.
func parseFlags(args []string) (configPath string, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild bool, rest []string) {
	rest = args[:0]
	for i := 0; i < len(args); i++ {
		switch {
		case args[i] == "-config" && i+1 < len(args):
			configPath = args[i+1]
			i++
		case strings.HasPrefix(args[i], "-config="):
			configPath = strings.TrimPrefix(args[i], "-config=")
		case args[i] == "--install-service" || args[i] == "-install-service":
			install = true
		case args[i] == "--remove-service" || args[i] == "-remove-service":
			remove = true
		case args[i] == "--no-copy" || args[i] == "-no-copy":
			noCopy = true
		case args[i] == "-h" || args[i] == "--help" || args[i] == "-help":
			help = true
		case args[i] == "-v" || args[i] == "--version" || args[i] == "-version":
			showVersion = true
		case args[i] == "--force-update" || args[i] == "-force-update":
			forceUpdate = true
		case args[i] == "--elevated-child":
			elevatedChild = true
		default:
			rest = append(rest, args[i])
		}
	}
	return configPath, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild, rest
}

// versionLine is printed at the top of every help and error screen.
func versionLine() string { return "gpu-turnstile " + version }

const usageText = `GPU arbitration proxy for Ollama + ComfyUI

Usage:
  gpu-turnstile -config <path>                                run the proxy
  gpu-turnstile --install-service [--no-copy] [-config path]  install + start as a service
  gpu-turnstile --remove-service                              stop + uninstall the service
  gpu-turnstile -v | --version                                print just the version
  gpu-turnstile --force-update                                check for a signed update now,
                                                              apply it and restart the service
  gpu-turnstile -h | --help                                   this help

Options:
  -config <path>      config file (default: gpu-turnstile.env next to the exe)
  --install-service   copies the binary into the canonical location
                      (%ProgramFiles%\gpu-turnstile or /var/lib/gpu-turnstile)
                      unless --no-copy; on Windows a UAC prompt appears when
                      the shell is not elevated
  --no-copy           with --install-service: register the current location as-is

All runtime settings are environment variables or KEY=VALUE lines in the
config file (OLLAMA_URL, COMFY_URL, LOGLEVEL, ...); see README.md.
`

// printHelp prints the version header plus the full help text.
func printHelp() {
	fmt.Printf("%s — %s", versionLine(), usageText)
}

// fatalUsage prints the version header, an error message and the one-line
// usage summary, then exits with code 2.
func fatalUsage(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "%s\n\n", versionLine())
	fmt.Fprintf(os.Stderr, format+"\n\n", args...)
	fmt.Fprintln(os.Stderr, "usage: gpu-turnstile [-config path] [--install-service [--no-copy] | --remove-service]")
	fmt.Fprintln(os.Stderr, "       gpu-turnstile service install|remove [-config path]")
	os.Exit(2)
}

func main() {
	configPath, install, remove, noCopy, help, showVersion, forceUpdate, elevatedChild, args := parseFlags(os.Args[1:])
	if showVersion {
		fmt.Println(version)
		return
	}
	bare := configPath == "" && !install && !remove && !forceUpdate && !elevatedChild && len(args) == 0
	if help || (bare && stdoutIsTerminal()) {
		// Bare invocation in a terminal (e.g. double-clicked on Windows)
		// shows the help instead of starting a proxy window with no visible
		// explanation. Without a terminal — Docker containers, services,
		// pipes — a bare invocation starts the proxy as before.
		printHelp()
		return
	}
	if len(args) > 0 && args[0] == "service" {
		// Legacy subcommand form: gpu-turnstile service install|remove.
		if len(args) != 2 || (args[1] != "install" && args[1] != "remove") {
			fatalUsage("error: expected 'service install' or 'service remove'")
		}
		install = args[1] == "install"
		remove = !install
		args = nil
	}
	switch {
	case install && remove:
		fatalUsage("error: --install-service and --remove-service are mutually exclusive")
	case forceUpdate && (install || remove):
		fatalUsage("error: --force-update cannot be combined with --install-service/--remove-service")
	case install:
		os.Exit(serviceCommand(configPath, true, noCopy, elevatedChild))
	case remove:
		os.Exit(serviceCommand(configPath, false, noCopy, elevatedChild))
	case forceUpdate:
		os.Exit(forceUpdateCommand(configPath, elevatedChild))
	}
	if len(args) > 0 {
		fatalUsage("error: unknown arguments: %s", strings.Join(args, " "))
	}

	if exePath, err := os.Executable(); err == nil {
		update.CleanupOld(exePath)
	}

	cfg, err := loadMergedConfig(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: %v\n", versionLine(), err)
		os.Exit(1)
	}
	log, logOut, logCloser := newLogger(cfg)
	defer logCloser.Close()
	syncEnvFile(resolveConfigPath(configPath), cfg.LogFile, log)

	if service.IsService() {
		if err := service.Run(func(ctx context.Context) error { return run(ctx, cfg, log, logOut, true) }); err != nil {
			log.Error("service failed", "err", err)
			os.Exit(1)
		}
		return
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	if err := run(ctx, cfg, log, logOut, false); err != nil {
		log.Error("listener failed", "err", err)
		os.Exit(1)
	}
}

// syncEnvFile upgrades an installer-written config file after an update:
// settings added since its CFG_VER are appended (commented out) and CFG_VER
// is bumped. Files not written by the installer (no CFG_VER), up-to-date
// files and dev builds are left untouched; a write failure is logged, not
// fatal.
func syncEnvFile(path, logFile string, log *slog.Logger) {
	data, err := os.ReadFile(path)
	if err != nil {
		return // no config file; nothing to upgrade
	}
	synced, changed := config.SyncSample(string(data), version, logFile)
	if !changed {
		return
	}
	if err := os.WriteFile(path, []byte(synced), 0o644); err != nil {
		log.Warn("could not append new settings to the config file", "path", path, "err", err)
		return
	}
	log.Warn("config file updated: new settings appended", "path", path, "version", version)
}

// defaultConfigPath returns gpu-turnstile.env next to the executable.
func defaultConfigPath() string {
	exe, err := os.Executable()
	if err != nil {
		return "gpu-turnstile.env"
	}
	return filepath.Join(filepath.Dir(exe), "gpu-turnstile.env")
}

// resolveConfigPath applies the precedence: -config flag, then
// GPU_TURNSTILE_CONFIG, then the default next to the executable.
func resolveConfigPath(flagValue string) string {
	if flagValue != "" {
		return flagValue
	}
	if v := os.Getenv("GPU_TURNSTILE_CONFIG"); v != "" {
		return v
	}
	return defaultConfigPath()
}

// loadMergedConfig reads the config file (if present) and overlays process
// environment variables on top. A missing file is fine; an unreadable or
// malformed file is fatal.
func loadMergedConfig(flagValue string) (config.Config, error) {
	path := resolveConfigPath(flagValue)
	values := map[string]string{}
	if f, err := os.Open(path); err == nil {
		defer f.Close()
		parsed, err := config.ParseEnvFile(f)
		if err != nil {
			return config.Config{}, fmt.Errorf("%s: %w", path, err)
		}
		values = parsed
	} else if !errors.Is(err, os.ErrNotExist) {
		return config.Config{}, fmt.Errorf("read config file: %w", err)
	}
	getenv := func(key string) string {
		if v := os.Getenv(key); v != "" {
			return v
		}
		return values[key]
	}
	return config.Load(getenv)
}

// newLogger builds the slog logger and returns the output writer (stderr or
// the opened LOG_FILE) plus a closer for it.
func newLogger(cfg config.Config) (*slog.Logger, io.Writer, io.Closer) {
	out := io.Writer(os.Stderr)
	closer := io.NopCloser(nil)
	if cfg.LogFile != "" {
		if f, err := os.OpenFile(cfg.LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644); err == nil {
			out = f
			closer = f
		} else {
			fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot open LOG_FILE %s: %v (logging to stderr)\n", cfg.LogFile, err)
		}
	}
	opts := &slog.HandlerOptions{Level: cfg.LogLevel}
	var handler slog.Handler = slog.NewTextHandler(out, opts)
	if cfg.LogJSON {
		handler = slog.NewJSONHandler(out, opts)
	}
	log := slog.New(handler)
	slog.SetDefault(log)
	return log, out, closer
}

// waitForEnter keeps an elevated child's console window open until the
// user has read the output.
func waitForEnter() {
	fmt.Print("\nPress Enter to close this window...")
	bufio.NewReader(os.Stdin).ReadString('\n')
}

// elevateAndMirror relaunches the current command elevated (UAC) and
// mirrors the child's exit code. verb is used in messages.
func elevateAndMirror(verb string) (int, bool) {
	args := append(append([]string{}, os.Args[1:]...), "--elevated-child")
	code, err := service.RelaunchElevated(args)
	if errors.Is(err, service.ErrUserCancelled) {
		fmt.Fprintln(os.Stderr, "gpu-turnstile: UAC prompt declined")
		return 1, true
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: could not elevate: %v\n", err)
		return 1, true
	}
	if code != 0 {
		fmt.Fprintf(os.Stderr, "gpu-turnstile %s failed in the elevated process (exit %d)\n", verb, code)
		return code, true
	}
	return 0, true
}

// isPermission reports whether err is a permission problem (Windows
// ERROR_ACCESS_DENIED, POSIX EACCES/EPERM, possibly wrapped).
func isPermission(err error) bool {
	return errors.Is(err, fs.ErrPermission) || strings.Contains(strings.ToLower(err.Error()), "access is denied")
}

// forceUpdateCommand checks for a signed update immediately, stages it if
// newer, and restarts the service when it is running so the new binary
// takes effect. Staging into a system directory and restarting a service
// need admin rights; instead of prompting unconditionally, permission
// failures trigger the UAC relaunch so a dev copy in a user-writable
// directory updates without a prompt.
func forceUpdateCommand(configPath string, elevatedChild bool) int {
	if elevatedChild {
		defer waitForEnter()
	}
	cfg, err := loadMergedConfig(configPath)
	if errors.Is(err, config.ErrNoConsumer) {
		err = nil // update settings do not depend on a consumer URL
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: %v\n", versionLine(), err)
		return 1
	}
	exePath, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: cannot locate executable: %v\n", err)
		return 1
	}
	log, _, logCloser := newLogger(cfg)
	defer logCloser.Close()
	if cfg.AppVersion == "dev" {
		fmt.Printf("%s: APP_VER=dev, updates disabled\n", versionLine())
		return 0
	}
	u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}

	// Single-shot: one attempt, fail fast when the server is unreachable
	// instead of hanging in a TCP connect for minutes.
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	staged, err := u.Check(ctx, exePath)
	if err != nil && isPermission(err) && !service.Elevated() {
		code, _ := elevateAndMirror("--force-update")
		if code == 0 {
			fmt.Println("update applied (elevated)")
		}
		return code
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n\ngpu-turnstile: update check failed: %v\n", versionLine(), err)
		return 1
	}
	if !staged {
		fmt.Printf("%s is up to date\n", versionLine())
		return 0
	}
	fmt.Printf("%s: update staged\n", versionLine())
	restarted, err := service.RestartIfRunning()
	if err != nil && isPermission(err) && !service.Elevated() {
		code, _ := elevateAndMirror("--force-update")
		if code == 0 {
			fmt.Println("update applied (elevated)")
		}
		return code
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile: update staged but service restart failed: %v\n", err)
		return 1
	}
	if restarted {
		fmt.Println("service restarted on the new version")
	} else {
		fmt.Println("no running service; the new version applies on next start")
	}
	return 0
}

// serviceCommand installs (copyBin = register the canonical-layout copy)
// or removes the service and reports the result. On Windows, when the
// shell is not elevated, the command relaunches itself through a UAC
// prompt and mirrors the elevated child's exit code. An elevated child
// waits for a keypress so its console window does not flash closed before
// the output can be read.
func serviceCommand(configPath string, install, noCopy, elevatedChild bool) int {
	verb, doneVerb := "remove", "removed"
	if install {
		verb, doneVerb = "install", "installed"
	}
	if elevatedChild {
		defer waitForEnter()
	}
	if !service.Elevated() {
		code, done := elevateAndMirror(verb)
		if done && code != 0 {
			return code
		}
		if done {
			fmt.Printf("service %s: %s (elevated)\n", service.Name, doneVerb)
			return 0
		}
	}
	var err error
	if install {
		path := resolveConfigPath(configPath)
		if abs, absErr := filepath.Abs(path); absErr == nil {
			path = abs
		}
		err = service.Install(path, !noCopy, version)
	} else {
		err = service.Remove()
	}
	if err != nil {
		fmt.Fprintf(os.Stderr, "gpu-turnstile service %s: %v\n", verb, err)
		return 1
	}
	fmt.Printf("service %s: %s\n", service.Name, doneVerb)
	return 0
}

// orDisabled renders an empty URL as "disabled" for the startup dump.
func orDisabled(url string) string {
	if url == "" {
		return "disabled"
	}
	return url
}

// managedComfyCommand resolves how ComfyUI is launched when it is managed:
// COMFY_CMD verbatim, or the standard venv layout under COMFY_DIR. Empty
// when neither is set (unmanaged).
func managedComfyCommand(cfg config.Config) string {
	if cfg.ComfyCmd != "" {
		return cfg.ComfyCmd
	}
	if cfg.ComfyDir != "" {
		return supervise.DefaultComfyCommand(runtime.GOOS, cfg.ComfyDir, cfg.ComfyURL)
	}
	return ""
}

func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Writer, isService bool) error {
	// ComfyUI can run as a managed child — COMFY_CMD verbatim, or the
	// standard venv layout derived from COMFY_DIR alone: started on demand
	// by the proxy, stopped after COMFY_IDLE_TIMEOUT idle (and on shutdown)
	// so its VRAM is freed.
	comfyCmdLine := managedComfyCommand(cfg)

	// The startup line carries the version and every setting and is emitted
	// at WARN so it is visible even with the default (quiet) log level.
	log.Log(ctx, slog.LevelWarn, "starting gpu-turnstile",
		"version", version,
		"listen_ollama", cfg.ListenOllama,
		"listen_comfy", cfg.ListenComfy,
		"ollama_url", orDisabled(cfg.OllamaURL),
		"comfy_url", orDisabled(cfg.ComfyURL),
		"unload_timeout", cfg.UnloadTimeout,
		"job_timeout", cfg.JobTimeout,
		"llm_wait_timeout", cfg.LLMWaitTimeout,
		"llm_busy_mode", cfg.LLMBusyMode,
		"llm_busy_status", cfg.LLMBusyStatus,
		"busy_retry_after", cfg.BusyRetryAfter,
		"unload_poll_interval", cfg.UnloadPollInterval,
		"history_poll_interval", cfg.HistoryPollInterval,
		"probe_timeout", cfg.ProbeTimeout,
		"health_interval", cfg.HealthInterval,
		"free_timeout", cfg.FreeTimeout,
		"warm_timeout", cfg.WarmTimeout,
		"shutdown_timeout", cfg.ShutdownTimeout,
		"backoff_initial", cfg.BackoffInitial,
		"backoff_max", cfg.BackoffMax,
		"prompt_capture_limit", cfg.PromptCaptureLimit,
		"warm_model", cfg.WarmModel,
		"comfy_cmd", orDisabled(comfyCmdLine),
		"comfy_dir", cfg.ComfyDir,
		"comfy_idle_timeout", cfg.ComfyIdleTimeout,
		"comfy_start_timeout", cfg.ComfyStartTimeout,
		"game_procs", cfg.GameProcs,
		"gpu_foreign_vram_mb", cfg.GPUForeignVRAMMB,
		"gpu_ignore_procs", cfg.GPUIgnoreProcs,
		"game_poll_interval", cfg.GamePollInterval,
		"auto_update", cfg.AutoUpdate,
		"update_interval", cfg.UpdateInterval,
		"update_repo", cfg.UpdateRepo,
		"update_asset", cfg.UpdateAsset,
		"log_level", cfg.LogLevel,
		"log_format", map[bool]string{true: "json", false: "text"}[cfg.LogJSON],
		"log_file", cfg.LogFile,
	)

	lk := lock.New(log)
	// Each consumer is enabled by setting its URL; a disabled consumer gets
	// no client, no listener and no probe.
	var ollamaClient *ollama.Client
	var err error
	if cfg.OllamaURL != "" {
		if ollamaClient, err = ollama.New(cfg.OllamaURL, log); err != nil {
			return err
		}
	}
	var comfyClient *comfy.Client
	if cfg.ComfyURL != "" {
		if comfyClient, err = comfy.New(cfg.ComfyURL, log); err != nil {
			return err
		}
	}

	var comfySup *supervise.Process
	if comfyCmdLine != "" {
		if cfg.ComfyCmd == "" {
			// Derived from COMFY_DIR: flag a wrong-looking layout early,
			// while the operator is still watching the startup log.
			python, script := supervise.ComfyLayout(runtime.GOOS, cfg.ComfyDir)
			for _, p := range []string{python, script} {
				if _, err := os.Stat(p); err != nil {
					log.Warn("COMFY_DIR: file not found; ComfyUI requests will fail until it exists", "path", p)
				}
			}
		}
		var err error
		comfySup, err = supervise.New("comfy", comfyCmdLine, cfg.ComfyDir, comfyClient.Probe, cfg.ComfyStartTimeout, log)
		if err != nil {
			return err
		}
		defer comfySup.Stop()
		gpuIdle := func() bool {
			state, _, pending := lk.Snapshot()
			return state == lock.StateIdle && !pending
		}
		go comfySup.WatchIdle(ctx, cfg.ComfyIdleTimeout, gpuIdle)
	}

	srv, err := proxy.New(proxy.Config{
		OllamaURL:           cfg.OllamaURL,
		ComfyURL:            cfg.ComfyURL,
		Lock:                lk,
		Ollama:              ollamaClient,
		Comfy:               comfyClient,
		ComfySup:            comfySup,
		Metrics:             metrics.New(),
		Log:                 log,
		LogColor:            !cfg.LogJSON && cfg.LogFile == "" && os.Getenv("NO_COLOR") == "",
		LogWriter:           logOut,
		LLMWaitTimeout:      cfg.LLMWaitTimeout,
		UnloadTimeout:       cfg.UnloadTimeout,
		JobTimeout:          cfg.JobTimeout,
		UnloadPollInterval:  cfg.UnloadPollInterval,
		HistoryPollInterval: cfg.HistoryPollInterval,
		FreeTimeout:         cfg.FreeTimeout,
		WarmTimeout:         cfg.WarmTimeout,
		LLMBusyMode:         cfg.LLMBusyMode,
		LLMBusyStatus:       cfg.LLMBusyStatus,
		BusyRetryAfter:      cfg.BusyRetryAfter,
		BackoffInitial:      cfg.BackoffInitial,
		BackoffMax:          cfg.BackoffMax,
		PromptCaptureLimit:  cfg.PromptCaptureLimit,
		WarmModel:           cfg.WarmModel,
	})
	if err != nil {
		return err
	}

	// Probe the enabled upstreams once; failure is logged, not fatal. A
	// managed ComfyUI is intentionally down at startup — the first request
	// starts it — so neither the startup probe nor the health check treats
	// that as an outage.
	probes := map[string]func(context.Context) error{}
	if ollamaClient != nil {
		probes["ollama"] = ollamaClient.Probe
	}
	if comfyClient != nil {
		if comfySup == nil {
			probes["comfy"] = comfyClient.Probe
		} else {
			// Managed upstream: an idle-stopped or still-starting server is
			// not an outage, so it is skipped until it has answered once
			// (Ready resets on every spawn/stop). After that, a failed
			// probe while the process lives is a real "DOWN".
			probes["comfy"] = func(ctx context.Context) error {
				if !comfySup.Ready() {
					if comfySup.Running() {
						if err := comfyClient.Probe(ctx); err == nil {
							comfySup.MarkReady()
						}
					}
					return errManagedDown
				}
				return comfyClient.Probe(ctx)
			}
		}
	}
	probeCtx, probeCancel := context.WithTimeout(ctx, cfg.ProbeTimeout)
	for name, probe := range probes {
		if err := probe(probeCtx); err != nil && !errors.Is(err, errManagedDown) {
			log.Warn(name+" probe failed", "err", err)
		}
	}
	probeCancel()
	if cfg.HealthInterval > 0 {
		go healthLoop(ctx, cfg.HealthInterval, cfg.ProbeTimeout, log, probes)
	}

	// Foreign GPU holders (games, other ML jobs) — enabled by GAME_PROCS
	// and/or GPU_FOREIGN_VRAM_MB — hold the lock externally while they run.
	if len(cfg.GameProcs) > 0 || cfg.GPUForeignVRAMMB > 0 {
		det := game.New(cfg.GameProcs, cfg.GPUForeignVRAMMB, cfg.GPUIgnoreProcs, log)
		go gameLoop(ctx, cfg, log, det, lk, ollamaClient, comfySup)
	}

	// Bind the listeners up front so a port conflict fails fast and the
	// readiness notification below really means "accepting connections".
	var servers []*http.Server
	var listeners []net.Listener
	bind := func(addr string, handler http.Handler, consumer string) error {
		ln, err := net.Listen("tcp", addr)
		if err != nil {
			return fmt.Errorf("listen %s on %s: %w", consumer, addr, err)
		}
		servers = append(servers, &http.Server{Addr: addr, Handler: handler})
		listeners = append(listeners, ln)
		log.Warn("listening", "consumer", consumer, "addr", addr)
		return nil
	}
	if ollamaClient != nil {
		if err := bind(cfg.ListenOllama, srv.OllamaHandler(), "ollama"); err != nil {
			return err
		}
	}
	if comfyClient != nil {
		if err := bind(cfg.ListenComfy, srv.ComfyHandler(), "comfy"); err != nil {
			return err
		}
	}

	errCh := make(chan error, len(servers))
	for i := range servers {
		go func(s *http.Server, ln net.Listener) { errCh <- s.Serve(ln) }(servers[i], listeners[i])
	}

	// Tell systemd we are up and start the watchdog pings; both are no-ops
	// when not running under a notify/watchdog unit.
	service.NotifyReady()
	service.StartWatchdog(ctx)

	if cfg.AutoUpdate {
		go updateLoop(ctx, cfg, log, lk, isService)
	}

	select {
	case err := <-errCh:
		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			return err
		}
	case <-ctx.Done():
		log.Info("shutting down")
	}

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
	defer shutdownCancel()
	for _, s := range servers {
		s.Shutdown(shutdownCtx)
	}
	return nil
}

// errManagedDown marks a managed upstream that is intentionally stopped
// (idle); health checks skip it instead of logging an outage.
var errManagedDown = errors.New("managed upstream intentionally stopped")

// gameLoop polls for foreign GPU holders (a game, another ML job). While one
// is detected it holds the lock externally so new LLM and image requests
// wait (or are rejected per LLM_BUSY_MODE), and — once in-flight work has
// drained — frees VRAM for it: the managed ComfyUI is stopped and Ollama's
// resident models are unloaded.
func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *game.Detector, lk *lock.Lock, ollamaClient *ollama.Client, comfySup *supervise.Process) {
	ticker := time.NewTicker(cfg.GamePollInterval)
	defer ticker.Stop()
	held, freed := false, false
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		holders, err := det.Check(ctx)
		if err != nil && ctx.Err() == nil {
			log.Warn("game detection failed", "err", err)
		}
		switch {
		case len(holders) > 0 && !held:
			held = true
			lk.SetExternal(summarizeHolders(holders))
			log.Warn("GPU held by an external process; new LLM/image requests wait",
				"holders", summarizeHolders(holders))
		case len(holders) == 0 && held:
			held, freed = false, false
			lk.ClearExternal()
			log.Warn("external process released the GPU; resuming")
		}
		if held && !freed {
			if state, _, _ := lk.Snapshot(); state != lock.StateLLM && state != lock.StateImage {
				freed = true
				freeVRAM(ctx, cfg.UnloadTimeout, log, ollamaClient, comfySup)
			}
		}
	}
}

// summarizeHolders joins holder descriptions for logs and busy responses,
// capping the list so a process name matching dozens of PIDs (system
// services) does not flood the log.
func summarizeHolders(holders []string) string {
	const max = 5
	if len(holders) > max {
		return strings.Join(holders[:max], "; ") + fmt.Sprintf("; +%d more", len(holders)-max)
	}
	return strings.Join(holders, "; ")
}

// freeVRAM stops the managed ComfyUI (never an external server on its port)
// and unloads Ollama's resident models so the foreign process gets the GPU
// memory.
func freeVRAM(ctx context.Context, unloadTimeout time.Duration, log *slog.Logger, ollamaClient *ollama.Client, comfySup *supervise.Process) {
	if comfySup != nil && comfySup.Running() {
		log.Warn("stopping the managed ComfyUI to free VRAM")
		comfySup.Stop()
	}
	if ollamaClient == nil {
		return
	}
	uctx, cancel := context.WithTimeout(ctx, unloadTimeout)
	defer cancel()
	if models, err := ollamaClient.LoadedModels(uctx); err != nil || len(models) == 0 {
		return // nothing resident (or ollama unreachable); nothing to free
	}
	elapsed, err := ollamaClient.UnloadAll(uctx)
	if err != nil {
		log.Warn("ollama unload incomplete; continuing", "err", err)
		return
	}
	log.Warn("ollama models unloaded to free VRAM", "seconds", elapsed.Seconds())
}

// healthLoop probes the enabled upstreams every interval and logs status
// transitions — "is DOWN" when a previously healthy upstream stops
// answering, "recovered" when it comes back. The first round only
// establishes the baseline; the startup probe already reported that state.
func healthLoop(ctx context.Context, interval, probeTimeout time.Duration, log *slog.Logger, probes map[string]func(context.Context) error) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	up := map[string]bool{}
	managed := map[string]bool{}
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		for name, probe := range probes {
			pctx, cancel := context.WithTimeout(ctx, probeTimeout)
			err := probe(pctx)
			cancel()
			if errors.Is(err, errManagedDown) {
				managed[name] = true // intentionally stopped; not an outage
				continue
			}
			now := err == nil
			if managed[name] {
				// First real probe after an idle stop only re-baselines —
				// an on-demand start is not a "recovery".
				managed[name] = false
				up[name] = now
				continue
			}
			was, seen := up[name]
			if seen && now != was {
				if now {
					log.Warn(name + " upstream recovered")
				} else {
					log.Warn(name+" upstream is DOWN", "err", err)
				}
			}
			up[name] = now
		}
	}
}

// updateLoop checks for signed updates on startup and every UPDATE_INTERVAL.
// In service mode a staged update is applied by exiting with exitCodeUpdate
// once the GPU lock is idle; the service recovery configuration restarts the
// process with the new binary. Interactively it only logs.
func updateLoop(ctx context.Context, cfg config.Config, log *slog.Logger, lk *lock.Lock, isService bool) {
	exePath, err := os.Executable()
	if err != nil {
		log.Warn("auto-update disabled: cannot locate executable", "err", err)
		return
	}
	u := &update.Updater{Repo: cfg.UpdateRepo, Asset: cfg.UpdateAsset, Version: version, Desired: cfg.AppVersion, Log: log}
	for {
		staged, err := u.Check(ctx, exePath)
		if err != nil && ctx.Err() == nil {
			log.Warn("auto-update check failed", "err", err)
		}
		if staged {
			if !isService {
				log.Warn("auto-update: new binary staged; restart gpu-turnstile to apply")
				return
			}
			log.Warn("auto-update: staged; restarting once the GPU is idle")
			if waitForIdle(ctx, lk, 24*time.Hour) {
				log.Warn("auto-update: restarting to apply update")
				os.Exit(exitCodeUpdate)
			}
			return
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(cfg.UpdateInterval):
		}
	}
}

// waitForIdle polls the lock until no LLM or image work is active or
// pending, max at most. Returns false on timeout or cancellation.
func waitForIdle(ctx context.Context, lk *lock.Lock, max time.Duration) bool {
	deadline := time.Now().Add(max)
	for {
		if state, _, pending := lk.Snapshot(); state == lock.StateIdle && !pending {
			return true
		}
		if time.Now().After(deadline) {
			return false
		}
		select {
		case <-ctx.Done():
			return false
		case <-time.After(5 * time.Second):
		}
	}
}