#!/bin/bash
# Shared helpers for the installer steps. Sourced by every step; not meant to be run alone.
# shellcheck disable=SC2034  # the settings below are used by the scripts that source this file

say()  { printf '\033[1m[wxch]\033[0m %s\n' "$*"; }
warn() { printf '\033[33m[wxch] warning:\033[0m %s\n' "$*" >&2; }
die()  { printf '\033[31m[wxch] error:\033[0m %s\n' "$*" >&2; exit 1; }
have() { command -v "$1" >/dev/null 2>&1; }

# ---- settings (all can be overridden with WXCH_* environment variables) ----
BASE="${WXCH_BASE:?WXCH_BASE is not set (run the installer through /get)}"
OWNER="${WXCH_USER:-${SUDO_USER:-}}"
[ -n "$OWNER" ] && [ "$OWNER" != root ] || die "cannot tell which user should own the agent; run through sudo from that user or set WXCH_USER"
OWNER_HOME="$(getent passwd "$OWNER" | cut -d: -f6)"
OWNER_UID="$(id -u "$OWNER")"
MODELS_DIR="${WXCH_MODELS_DIR:-/mnt/models}"
# external disks (added with add-disk) are mounted on demand at $EXT_ROOT/<id> and used for downloaded models
EXT_ROOT="${WXCH_EXT_ROOT:-/mnt/wxch-ext}"
AGENT_DIR=/opt/wxch-agent
AGENT_BIN="$AGENT_DIR/wxch-agent"
VLLM_IMAGE="${WXCH_VLLM_IMAGE:-nvcr.io/nvidia/vllm:26.09-py3}"

# the smallest model of the same family on each runtime, enough to prove that everything works
MODEL_OLLAMA="${WXCH_MODEL_OLLAMA:-qwen2.5:0.5b}"
MODEL_LMSTUDIO="${WXCH_MODEL_LMSTUDIO:-lmstudio-community/Qwen2.5-0.5B-Instruct-GGUF@Q4_K_M}"
MODEL_VLLM="${WXCH_MODEL_VLLM:-Qwen/Qwen2.5-0.5B-Instruct}"

case "$(uname -m)" in
  aarch64|arm64) ARCH=arm64 ;;
  x86_64|amd64)  ARCH=amd64 ;;
  *) die "unsupported CPU architecture $(uname -m)" ;;
esac

# run a command as the agent's user (a plain root->user switch, no password involved)
as_owner() { sudo -u "$OWNER" -H "$@"; }

# WXCH_SKIP="vllm,lmstudio" skips steps; WXCH_ONLY="vllm" runs just those
step_enabled() {
  name="$1"
  [ -e "/run/wxch-auto-skip-$name" ] && return 1   # preflight decided this machine cannot run it (e.g. vLLM without Docker/GPU)
  if [ -n "${WXCH_ONLY:-}" ]; then
    case ",$WXCH_ONLY," in *",$name,"*) return 0 ;; *) return 1 ;; esac
  fi
  case ",${WXCH_SKIP:-}," in *",$name,"*) return 1 ;; esac
  return 0
}

# install OS packages with whatever package manager this distribution has
pkg_install() {
  if   have apt-get; then DEBIAN_FRONTEND=noninteractive apt-get install -y -qq "$@" >/dev/null
  elif have dnf;     then dnf install -y -q "$@" >/dev/null
  elif have yum;     then yum install -y -q "$@" >/dev/null
  elif have zypper;  then zypper --non-interactive -q install "$@" >/dev/null
  elif have pacman;  then pacman -S --noconfirm --needed "$@" >/dev/null
  elif have apk;     then apk add -q "$@" >/dev/null
  else return 1; fi
}
# package names differ per family: print the right name for a generic one
pkg_name() {
  case "$1" in
    libatomic) if have apt-get; then echo libatomic1; elif have pacman || have apk; then echo libatomic_ops; else echo libatomic; fi ;;
    *) echo "$1" ;;
  esac
}
pkg_hint() { # command to show people when something is missing
  if   have apt-get; then echo "apt-get install $*"
  elif have dnf;     then echo "dnf install $*"
  elif have yum;     then echo "yum install $*"
  elif have zypper;  then echo "zypper install $*"
  elif have pacman;  then echo "pacman -S $*"
  else echo "install: $*"; fi
}
# SELinux (Fedora/RHEL family): files copied into place need the right label to be started by systemd
fix_selinux() { have restorecon && [ -e /sys/fs/selinux/enforce ] && restorecon -R "$@" >/dev/null 2>&1 || true; }

# download with retries; the file is only moved into place when complete
fetch() { # url dest
  url="$1"; dest="$2"; n=0
  until curl -fsSL --connect-timeout 20 --retry 3 --retry-delay 3 -o "$dest.part" "$url"; do
    n=$((n+1)); [ "$n" -lt 3 ] || die "could not download $url"
    sleep 3
  done
  mv "$dest.part" "$dest"
}

# merge keys into /etc/wxch/node.json (kept as JSON, other keys preserved)
node_json_set() { # key value [key value ...]
  mkdir -p /etc/wxch
  python3 - "$@" <<'PY'
import json, os, sys
p = "/etc/wxch/node.json"
d = json.load(open(p)) if os.path.exists(p) else {}
a = sys.argv[1:]
for i in range(0, len(a), 2):
    v = a[i + 1]
    d[a[i]] = json.loads(v) if v.startswith("[") else (int(v) if v.isdigit() else v)
json.dump(d, open(p, "w"), indent=2)
open(p, "a").write("\n")
PY
}

# ---- external disks: pieces shared by add-disk and the install steps (all safe to repeat, none restarts anything) ----

# the helper that prepares a disk once it is mounted (marker file and folders owned by the agent user), run by systemd
# at every mount and by add-disk right away
install_ext_prep() {
  mkdir -p /etc/wxch
  echo "$OWNER" > /etc/wxch/ext-owner
  cat > /usr/local/sbin/wxch-ext-prep <<'PREP'
#!/bin/sh
# usage: wxch-ext-prep <mount point> [agent user]   (idempotent; the user defaults to /etc/wxch/ext-owner)
EXT="$1"; OWNER="${2:-$(cat /etc/wxch/ext-owner 2>/dev/null)}"
[ -n "$EXT" ] && [ -n "$OWNER" ] || exit 1
mountpoint -q "$EXT" || exit 0
case "$(findmnt -no FSTYPE -T "$EXT" | tail -1)" in
  ext4|xfs|btrfs) ;;
  *) echo "wxch: the disk at $EXT is not ext4/xfs/btrfs: not used for models" >&2; exit 0 ;;
esac
S="$EXT/wxch-models"
for d in "$S" "$S/ollama" "$S/ollama/blobs" "$S/vllm"; do
  mkdir -p "$d"
  chown "$OWNER" "$d"
  chmod 755 "$d"
done
: > "$S/.wxch-models"
chown "$OWNER" "$S/.wxch-models"
PREP
  chmod 755 /usr/local/sbin/wxch-ext-prep
  cat > /etc/systemd/system/wxch-ext-prep@.service <<'UNIT'
[Unit]
Description=Prepare the wxch external model disk at %f
After=%i.mount
ConditionPathIsMountPoint=%f

[Service]
Type=oneshot
ExecStart=/usr/local/sbin/wxch-ext-prep %f
UNIT
  fix_selinux /usr/local/sbin/wxch-ext-prep /etc/systemd/system/wxch-ext-prep@.service
}

# The agent moves model blobs to and from external disks, so it needs to write in Ollama's blobs/ (and nowhere else in
# Ollama's store). An ACL grants exactly that to the agent's user; without ACL tools the user joins Ollama's group instead.
grant_ollama_blobs() {
  [ -d "$MODELS_DIR/ollama/blobs" ] || return 0
  have setfacl || pkg_install acl || true
  if have setfacl; then
    setfacl -m "u:$OWNER:rwx" "$MODELS_DIR/ollama/blobs" || warn "could not grant $OWNER access to $MODELS_DIR/ollama/blobs: Ollama models cannot be moved to an external disk"
  elif getent group ollama >/dev/null; then
    usermod -aG ollama "$OWNER" && chmod g+rwx "$MODELS_DIR/ollama/blobs" &&
      warn "$OWNER joined the ollama group: restart the agent (sudo systemctl restart wxch-agent, when no download is running) to move Ollama models to an external disk"
  fi
}

# let the vLLM runner start models that were moved to an external disk: it may follow links into this folder and mounts
# it into the container. The runner reads its config again for every request, so no restart (which would stop the
# running containers) is needed.
allow_runner_root() { # folder
  [ -f /etc/wxch/runner.json ] || return 0
  python3 - "$1" <<'PY'
import json, sys
p = "/etc/wxch/runner.json"
c = json.load(open(p))
roots = c.get("extra_roots", [])
if sys.argv[1] not in roots:
    roots.append(sys.argv[1])
    c["extra_roots"] = roots
    json.dump(c, open(p, "w"), indent=2)
    open(p, "a").write("\n")
PY
}