// wxch-runner: a small root service that lets the (unprivileged) agent start and stop vLLM containers
// without giving the agent Docker access. Docker access is equivalent to root, so this service only
// accepts a handful of operations and decides every security-relevant setting itself:
//   - the image, the models folder and the port range come from /etc/wxch/runner.json (owned by root)
//   - the model must be a Hugging Face style id with files that really exist under the models folder
//   - only an allow-list of vLLM flags is accepted, with validated values
//   - containers are never privileged, publish only on 127.0.0.1, and are limited in number
package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Config struct {
	Image         string   `json:"image"`
	ModelsDir     string   `json:"models_dir"`
	ModulesDir    string   `json:"modules_dir"` // engine folders (ComfyUI, Whisper) with their models
	AgentGID      int      `json:"agent_gid"`   // -1 = same as the uid
	ExtraRoots    []string `json:"extra_roots"` // more folders models may be linked to (the external disk), set by root
	AgentUID      int      `json:"agent_uid"`
	AgentExe      string   `json:"agent_exe"` // when set, a caller that runs as the agent's user must also be this program
	PortMin       int      `json:"port_min"`
	PortMax       int      `json:"port_max"`
	MaxContainers int      `json:"max_containers"`
	ShmSize       string   `json:"shm_size"`
	// TrustRemoteCode lists the exact models that may run the Python code shipped in their repository
	// (vLLM --trust-remote-code). Only root can fill it in (this file); the agent cannot ask for the flag.
	TrustRemoteCode []string `json:"trust_remote_code,omitempty"`
}

func defaultConfig() Config {
	return Config{Image: "nvcr.io/nvidia/vllm:26.09-py3", ModelsDir: "/mnt/models/vllm", ModulesDir: "/mnt/models/modules", AgentGID: -1, AgentUID: -1, PortMin: 8100, PortMax: 8199, MaxContainers: 2, ShmSize: "16g"}
}

const labelKey = "wxch.vllm"

var hfRepoRe = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$`)

func slugOf(name string) string {
	return strings.NewReplacer("/", "--", ".", "-", "_", "-").Replace(strings.ToLower(name))
}
func containerOf(name string) string { return "wxch-vllm-" + slugOf(name) }

// ---- request validation ----

type Request struct {
	Op     string   `json:"op"`
	Model  string   `json:"model,omitempty"`
	Port   int      `json:"port,omitempty"`
	Mem    int64    `json:"mem,omitempty"`
	Util   float64  `json:"util,omitempty"`
	MaxLen int      `json:"max_len,omitempty"`
	Args   []string `json:"args,omitempty"`
	Tail   int      `json:"tail,omitempty"`
	Engine string   `json:"engine,omitempty"`
}

type Response struct {
	OK   bool            `json:"ok"`
	Err  string          `json:"err,omitempty"`
	Data json.RawMessage `json:"data,omitempty"`
}

var (
	wordRe = regexp.MustCompile(`^[a-z0-9_]{1,40}$`)
	intRe  = regexp.MustCompile(`^[0-9]{1,5}$`)
)

// allowedFlags: flag -> validator for its value (nil = flag without value).
var allowedFlags = map[string]func(string) bool{
	"--kv-cache-dtype":          func(v string) bool { return v == "auto" || v == "fp8" || v == "fp8_e4m3" || v == "fp8_e5m2" },
	"--dtype":                   func(v string) bool { return v == "auto" || v == "bfloat16" || v == "float16" },
	"--enable-prefix-caching":   nil,
	"--enable-auto-tool-choice": nil,
	"--reasoning-parser":        wordRe.MatchString,
	"--tool-call-parser":        wordRe.MatchString,
	"--quantization":            wordRe.MatchString,
	"--max-num-seqs": func(v string) bool {
		n, err := strconv.Atoi(v)
		return intRe.MatchString(v) && err == nil && n >= 1 && n <= 1024
	},
}

// validateArgs accepts only allow-listed vLLM flags; returns them in order.
func validateArgs(args []string) ([]string, error) {
	var out []string
	for i := 0; i < len(args); i++ {
		a := args[i]
		check, ok := allowedFlags[a]
		if !ok {
			return nil, fmt.Errorf("vLLM option %q is not allowed", a)
		}
		out = append(out, a)
		if check != nil {
			if i+1 >= len(args) || !check(args[i+1]) {
				return nil, fmt.Errorf("invalid value for %s", a)
			}
			out = append(out, args[i+1])
			i++
		}
	}
	return out, nil
}

type container struct {
	Name  string `json:"name"`
	Model string `json:"model"`
	Port  int    `json:"port"`
	Mem   int64  `json:"mem"`
	State string `json:"state"`
}

// buildRun validates a start request and returns the exact docker arguments.
func buildRun(cfg Config, req Request, existing []container) ([]string, error) {
	if !hfRepoRe.MatchString(req.Model) {
		return nil, fmt.Errorf("invalid model name")
	}
	root, err := filepath.EvalSymlinks(cfg.ModelsDir)
	if err != nil {
		return nil, fmt.Errorf("models folder missing: %v", err)
	}
	dir, err := filepath.EvalSymlinks(filepath.Join(root, req.Model))
	if err != nil {
		return nil, fmt.Errorf("model files not found under the models folder")
	}
	mounts := []string{"-v", root + ":/models:ro"}
	if !strings.HasPrefix(dir, root+string(os.PathSeparator)) {
		// a model moved to the external disk is a link out of the models folder: allowed only into a folder that root
		// listed in the config, which is then mounted at the same path so the link resolves inside the container
		extraRoot := ""
		for _, e := range cfg.ExtraRoots {
			if er, err := filepath.EvalSymlinks(e); err == nil && strings.HasPrefix(dir, er+string(os.PathSeparator)) {
				extraRoot = er
				break
			}
		}
		if extraRoot == "" {
			return nil, fmt.Errorf("model files not found under the models folder")
		}
		mounts = append(mounts, "-v", extraRoot+":"+extraRoot+":ro")
	}
	if _, err := os.Stat(filepath.Join(dir, "config.json")); err != nil {
		return nil, fmt.Errorf("model folder has no config.json")
	}
	if req.Port < cfg.PortMin || req.Port > cfg.PortMax {
		return nil, fmt.Errorf("port outside the allowed range %d-%d", cfg.PortMin, cfg.PortMax)
	}
	running := 0
	for _, c := range existing {
		if c.Port == req.Port && c.Name != containerOf(req.Model) {
			return nil, fmt.Errorf("port %d is already used", req.Port)
		}
		if c.Name != containerOf(req.Model) {
			running++
		}
	}
	if running >= cfg.MaxContainers {
		return nil, fmt.Errorf("at most %d vLLM containers may exist; unload one first", cfg.MaxContainers)
	}
	if req.Util < 0.05 || req.Util > 0.9 {
		return nil, fmt.Errorf("gpu memory utilisation must be between 0.05 and 0.9")
	}
	if req.MaxLen < 512 || req.MaxLen > 1<<20 {
		return nil, fmt.Errorf("invalid max model length")
	}
	if req.Mem < 0 {
		return nil, fmt.Errorf("invalid memory figure")
	}
	extra, err := validateArgs(req.Args)
	if err != nil {
		return nil, err
	}
	args := []string{"run", "-d", "--name", containerOf(req.Model),
		"--label", labelKey + "=1", "--label", "wxch.model=" + req.Model, "--label", "wxch.port=" + strconv.Itoa(req.Port), "--label", "wxch.mem=" + strconv.FormatInt(req.Mem, 10),
		"--gpus", "all", "--shm-size", cfg.ShmSize, "--ulimit", "memlock=-1", "--ulimit", "stack=67108864",
		"--security-opt", "no-new-privileges", "--cap-drop", "ALL",
		"-e", "HF_HUB_OFFLINE=1", "-e", "VLLM_NO_USAGE_STATS=1", "-e", "DO_NOT_TRACK=1",
		"-p", fmt.Sprintf("127.0.0.1:%d:8000", req.Port)}
	args = append(args, mounts...)
	args = append(args, cfg.Image,
		"vllm", "serve", "/models/"+req.Model, "--served-model-name", req.Model, "--host", "0.0.0.0", "--port", "8000",
		"--max-model-len", strconv.Itoa(req.MaxLen), "--gpu-memory-utilization", fmt.Sprintf("%.2f", req.Util))
	for _, m := range cfg.TrustRemoteCode {
		if m == req.Model {
			args = append(args, "--trust-remote-code")
			break
		}
	}
	return append(args, extra...), nil
}

// ---- docker ----

func docker(timeout time.Duration, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, "docker", args...)
	var out, errb bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &errb
	if err := cmd.Run(); err != nil {
		msg := strings.TrimSpace(errb.String())
		if msg == "" {
			msg = err.Error()
		}
		return out.String(), fmt.Errorf("%s", msg)
	}
	return out.String(), nil
}

// dockerLogs returns what `docker logs` prints. It writes the container's stderr to its own stderr, and vLLM logs there, so
// reading stdout alone returned nothing and "vLLM exited while loading" came without a reason.
func dockerLogs(timeout time.Duration, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, "docker", args...)
	var out bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &out
	err := cmd.Run()
	return out.String(), err
}

func listContainers() ([]container, error) {
	out, err := docker(10*time.Second, "ps", "-a", "--filter", "label="+labelKey,
		"--format", `{{.Names}}|{{.Label "wxch.model"}}|{{.Label "wxch.port"}}|{{.Label "wxch.mem"}}|{{.State}}`)
	if err != nil {
		return nil, err
	}
	var cs []container
	for _, line := range strings.Split(strings.TrimSpace(out), "\n") {
		f := strings.Split(line, "|")
		if len(f) < 5 {
			continue
		}
		port, _ := strconv.Atoi(f[2])
		mem, _ := strconv.ParseInt(f[3], 10, 64)
		cs = append(cs, container{f[0], f[1], port, mem, f[4]})
	}
	return cs, nil
}

func ok(v any) Response {
	b, _ := json.Marshal(v)
	return Response{OK: true, Data: b}
}
func fail(format string, a ...any) Response { return Response{Err: fmt.Sprintf(format, a...)} }

// handle runs one request. The caller has already authenticated the peer.
func handle(cfg Config, req Request) Response {
	if strings.HasPrefix(req.Op, "mod_") {
		return handleModule(cfg, req)
	}
	switch req.Op {
	case "status":
		if _, err := docker(8*time.Second, "info", "--format", "{{.ServerVersion}}"); err != nil {
			return fail("docker: %v", err)
		}
		if _, err := docker(8*time.Second, "image", "inspect", "--format", "{{.Id}}", cfg.Image); err != nil {
			return fail("image %s is not pulled (run setup/40-vllm.sh)", cfg.Image)
		}
		return ok(map[string]any{"image": cfg.Image, "models_dir": cfg.ModelsDir, "max_containers": cfg.MaxContainers})
	case "ours":
		return ok(ourContainerIDs())
	case "ps":
		cs, err := listContainers()
		if err != nil {
			return fail("%v", err)
		}
		return ok(cs)
	case "start":
		cs, err := listContainers()
		if err != nil {
			return fail("%v", err)
		}
		args, err := buildRun(cfg, req, cs)
		if err != nil {
			return fail("%v", err)
		}
		docker(40*time.Second, "rm", "-f", containerOf(req.Model)) // a stale container of the same model
		if _, err := docker(5*time.Minute, args...); err != nil {
			return fail("%v", err)
		}
		return ok(map[string]any{"name": containerOf(req.Model)})
	case "stop":
		if !hfRepoRe.MatchString(req.Model) {
			return fail("invalid model name")
		}
		if _, err := docker(60*time.Second, "rm", "-f", containerOf(req.Model)); err != nil {
			return fail("%v", err)
		}
		return ok(nil)
	case "running":
		if !hfRepoRe.MatchString(req.Model) {
			return fail("invalid model name")
		}
		out, _ := docker(8*time.Second, "inspect", "-f", "{{.State.Running}}", containerOf(req.Model))
		return ok(strings.TrimSpace(out) == "true")
	case "logs":
		if !hfRepoRe.MatchString(req.Model) {
			return fail("invalid model name")
		}
		t := req.Tail
		if t < 1 || t > 200 {
			t = 30
		}
		out, err := dockerLogs(15*time.Second, "logs", "--tail", strconv.Itoa(t), containerOf(req.Model))
		if err != nil && out == "" {
			return fail("%v", err)
		}
		return ok(out)
	}
	return fail("unknown operation")
}

// ---- server ----

// readConfig loads the root-owned config. It is read again for every connection, so adding an external disk (a new
// extra_roots entry) needs no restart, which would stop the running vLLM containers.
func readConfig(cfgPath string) (Config, error) {
	cfg := defaultConfig()
	b, err := os.ReadFile(cfgPath)
	if err != nil {
		return cfg, fmt.Errorf("cannot read %s: %v", cfgPath, err)
	}
	if err := json.Unmarshal(b, &cfg); err != nil {
		return cfg, err
	}
	if cfg.AgentUID < 0 {
		return cfg, fmt.Errorf("agent_uid must be set in the config")
	}
	return cfg, nil
}

// peerAllowed says whether a caller may use the runner. Root always may. The agent's user may when the caller is the agent
// program itself (agent_exe in the config): any other program that runs as the same user, for example an AI engine started
// directly on the machine, is refused. Without agent_exe in the config (an older install) only the user is checked.
func peerAllowed(cfg Config, uid, pid int, exeOf func(int) (string, error)) bool {
	if uid == 0 {
		return true
	}
	if uid != cfg.AgentUID {
		return false
	}
	if cfg.AgentExe == "" {
		return true
	}
	exe, err := exeOf(pid)
	if err != nil {
		return false
	}
	want, werr := filepath.EvalSymlinks(cfg.AgentExe)
	if werr != nil {
		want = cfg.AgentExe
	}
	return filepath.Clean(exe) == filepath.Clean(want)
}

func serve(cfgPath, sockPath string) {
	cfg, err := readConfig(cfgPath)
	if err != nil {
		log.Fatal(err)
	}
	os.MkdirAll(filepath.Dir(sockPath), 0755)
	os.Remove(sockPath)
	l, err := net.Listen("unix", sockPath)
	if err != nil {
		log.Fatal(err)
	}
	os.Chmod(sockPath, 0666) // access is decided by the peer uid check below
	log.Printf("wxch-runner listening on %s (agent uid %d, image %s)", sockPath, cfg.AgentUID, cfg.Image)
	for {
		c, err := l.Accept()
		if err != nil {
			log.Fatal(err)
		}
		go func(c net.Conn) {
			defer c.Close()
			cur := cfg // the config of this connection: the file as it is now, or the last good one
			if fresh, err := readConfig(cfgPath); err == nil {
				cur = fresh
			} else {
				log.Printf("config not reloaded: %v", err)
			}
			uid, pid, err := peerCred(c)
			if err == nil && !peerAllowed(cur, uid, pid, peerExe) {
				err = fmt.Errorf("not the agent program")
			}
			if err != nil || (uid != cur.AgentUID && uid != 0) {
				json.NewEncoder(c).Encode(fail("not allowed"))
				log.Printf("refused connection from uid %d", uid)
				return
			}
			c.SetDeadline(time.Now().Add(10 * time.Minute))
			line, err := bufio.NewReaderSize(c, 64<<10).ReadBytes('\n')
			if err != nil {
				return
			}
			var req Request
			if json.Unmarshal(line, &req) != nil {
				json.NewEncoder(c).Encode(fail("bad request"))
				return
			}
			log.Printf("op %s %s%s", req.Op, req.Model, req.Engine)
			resp := handle(cur, req)
			if !resp.OK { // so that `journalctl -u wxch-runner` says why an operation was refused or failed
				log.Printf("op %s %s%s failed: %s", req.Op, req.Model, req.Engine, resp.Err)
			}
			json.NewEncoder(c).Encode(resp)
		}(c)
	}
}

func main() {
	cfgPath := flag.String("config", "/etc/wxch/runner.json", "config file (owned by root)")
	sock := flag.String("socket", "/run/wxch/runner.sock", "unix socket path")
	flag.Parse()
	serve(*cfgPath, *sock)
}