#!/bin/bash
# Optional AI modules (image, video, speech). NOT installed by default: this step does nothing unless WXCH_MODULES says
# which ones, as a comma separated list of capabilities or "all":
#
#   image.generate  image.edit  video.generate  video.edit  audio.transcribe
#
# The same can be done later from the control page (AI modules) or with:  curl -fsSL <server>/get | sh -s -- modules <names|all>
# The engines run directly on the machine, as the agent's own child processes: the agent installs a Python environment per
# engine (uv, Python 3.12, PyTorch for CUDA 13: about 5 GB, on the first install), downloads the engine's code and the model
# files. No Docker and no wxch-runner are needed. (A machine can still be switched to engines in Docker containers with
# "mode": "docker" in <modules folder>/state.json; then wxch-runner has to run.) Safe to re-run: what is installed is
# skipped, an interrupted download continues.
. "$(dirname "$0")/lib.sh"

list="${WXCH_MODULES:-}"
if [ -z "$list" ]; then
  say "no AI modules requested (they are optional); to add them later: curl -fsSL $BASE/get | sh -s -- modules <names|all>, or use the AI modules page"
  exit 0
fi
case "$list" in *[!a-z.,]*) die "WXCH_MODULES may only contain capability names (letters and dots) separated by commas" ;; esac

[ -x "$AGENT_BIN" ] || die "the agent is not installed ($AGENT_BIN); run the 'agent' step first"
# the runner is only needed when the machine is set to run engines in Docker
if grep -q '"mode"[[:space:]]*:[[:space:]]*"docker"' "${WXCH_MODELS_DIR:-/mnt/models}/modules/state.json" 2>/dev/null; then
  [ -S /run/wxch/runner.sock ] || die "this machine runs its engines in Docker, but wxch-runner is not running: the 'vllm' step has to run first"
fi
if ! out="$(as_owner "$AGENT_BIN" modules list 2>&1)"; then
  echo "$out" | head -5 >&2
  die "the agent program ($AGENT_BIN, version $("$AGENT_BIN" -version 2>&1 | head -1)) cannot list AI modules. Version 0.12.1 or newer is needed: run the full installer again (the 'agent' step updates it) or use Update agent on the Models page, then check the version with: $AGENT_BIN -version"
fi

ids="$(echo "$list" | tr ',' ' ')"
say "installing AI modules: $ids"
# shellcheck disable=SC2086  # the names are validated above and must be separate words
as_owner "$AGENT_BIN" modules install $ids || die "could not install the AI modules (see above); run the same command again to continue where it stopped"
say "AI modules are ready: they appear on the control page and on the studio page"