#!/bin/bash
# Prepare the vLLM runtime: the limited root helper (wxch-runner), the NVIDIA vLLM image and the model folder.
# Safe to re-run (an interrupted image pull continues with the layers already downloaded).
#
# SECURITY: the agent is NOT added to the docker group (that would be root without a password for anyone who can
# reach the agent). Instead wxch-runner, a small root service, starts/stops vLLM containers on its behalf. It
# accepts only a few operations and fixes the image, model folder and ports itself (/etc/wxch/runner.json).
. "$(dirname "$0")/lib.sh"

have docker || die "docker is not installed"
docker info >/dev/null 2>&1 || die "docker is not running"

mkdir -p "$MODELS_DIR/vllm"
chown "$OWNER" "$MODELS_DIR/vllm"
# optional AI modules (image, video, speech) are installed from the control page into this folder; empty until then
mkdir -p "$MODELS_DIR/modules"
chown "$OWNER" "$MODELS_DIR/modules"

say "installing wxch-runner ($ARCH)"
fetch "$BASE/dl/wxch-runner-$ARCH" /usr/local/bin/wxch-runner.new
chmod 755 /usr/local/bin/wxch-runner.new
mv /usr/local/bin/wxch-runner.new /usr/local/bin/wxch-runner

mkdir -p /etc/wxch
python3 - "$VLLM_IMAGE" "$MODELS_DIR/vllm" "$OWNER_UID" "$EXT_ROOT" "$MODELS_DIR/modules" "$(id -g "$OWNER")" "$AGENT_BIN" "${WXCH_TRUST_REMOTE_CODE:-}" <<'PY'
import json, re, sys
image, models_dir, uid, ext, modules_dir, gid, agent_exe, trust_new = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4], sys.argv[5], int(sys.argv[6]), sys.argv[7], sys.argv[8]
# models allowed to run the code shipped in their repository: kept across re-runs, added to only by the owner here
try:
    trust = list(json.load(open("/etc/wxch/runner.json")).get("trust_remote_code", []))
except Exception:
    trust = []
for m in [x.strip() for x in trust_new.split(",") if x.strip()]:
    if not re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", m):
        sys.exit("WXCH_TRUST_REMOTE_CODE: %r is not a Hugging Face id like org/name" % m)
    if m not in trust:
        trust.append(m)
# extra_roots: models moved to external disks are links into this folder, which is then mounted into the container
cfg = {"image": image, "models_dir": models_dir, "modules_dir": modules_dir, "agent_uid": uid, "agent_gid": gid, "agent_exe": agent_exe, "extra_roots": [ext],
       "port_min": 8100, "port_max": 8199, "max_containers": 2, "shm_size": "16g"}
if trust:
    cfg["trust_remote_code"] = trust
open("/etc/wxch/runner.json", "w").write(json.dumps(cfg, indent=2) + "\n")
PY
chown root:root /etc/wxch/runner.json
chmod 600 /etc/wxch/runner.json

cat > /etc/systemd/system/wxch-runner.service <<EOF
[Unit]
Description=wxch runner (limited root helper for vLLM containers)
After=docker.service network-online.target
Wants=docker.service
$(mountpoint -q "$MODELS_DIR" && echo "RequiresMountsFor=$MODELS_DIR")

[Service]
ExecStart=/usr/local/bin/wxch-runner -config /etc/wxch/runner.json
Restart=always
RestartSec=3
RuntimeDirectory=wxch
RuntimeDirectoryMode=0755
NoNewPrivileges=true
PrivateTmp=true

[Install]
WantedBy=multi-user.target
EOF
systemctl daemon-reload
systemctl enable wxch-runner >/dev/null 2>&1
systemctl restart wxch-runner

i=0
until [ -S /run/wxch/runner.sock ]; do
  i=$((i+1)); [ "$i" -lt 20 ] || die "the runner did not start; see: journalctl -u wxch-runner -n 30"
  sleep 1
done
say "wxch-runner is running"

if docker image inspect "$VLLM_IMAGE" >/dev/null 2>&1; then
  say "image $VLLM_IMAGE is already present"
else
  say "pulling $VLLM_IMAGE (about 11 GB; on a slow connection this takes a long time, and it can be resumed by running the same command again)"
  # A flaky connection (TLS handshake timeout, reset) is common for an 11 GB pull. Layers already downloaded are
  # kept, so retrying continues where it stopped.
  n=0
  until docker pull "$VLLM_IMAGE"; do
    n=$((n + 1))
    [ "$n" -lt 10 ] || die "docker pull failed after $n attempts: check the internet connection, then run the same command again (finished layers are kept)"
    say "pull interrupted (attempt $n of 10), retrying in $((n * 15)) s"
    sleep $((n * 15))
  done
fi
say "vLLM runtime ready"