#!/usr/bin/env bash
# Almond Axol installer.
#
#   curl https://axol.almond.bot/install -fsS | bash
#
# Installs uv, installs the axol CLI as a uv tool (from GitHub, pinned to the
# latest release tag, all extras), installs the pyzed bindings when the ZED SDK
# is present, installs adb for Quest-over-USB teleop, and registers a root
# systemd service that keeps `axol serve` running at boot and restarts it if it
# ever goes down. Safe to re-run: re-running upgrades to the latest release.

set -euo pipefail

INSTALL_URL="https://axol.almond.bot/install"
# GitHub stays the source of truth for what the latest release is; the package
# itself installs from PyPI (published by the release workflow), which is a
# prebuilt wheel — no repo clone or local build.
REPO_URL="https://github.com/almond-bot/axol"
# `video` is intentionally omitted: headset-video support is system-level
# (GStreamer NVENC + PyGObject + the patched zed-gstreamer source plugins),
# provisioned below by `axol provision`, not via a PyPI extra (see
# pyproject.toml). WebRTC transport itself is aiortc, a normal core dependency.
EXTRAS="lerobot,sim"
PYTHON_VERSION="3.13"
BIN_DIR="/usr/local/bin"
SERVICE_NAME="axol"
SERVICE_FILE="/etc/systemd/system/${SERVICE_NAME}.service"
# We install as root (the systemd service runs as root), but operators log in
# as their own user and expect `axol` to work there too. uv's default tool +
# managed-Python locations live under /root (mode 0700), which only root can
# traverse, so the `axol` wrapper on PATH would fail for everyone else. Pin
# both to shared, world-readable locations instead. The same vars are baked
# into the systemd unit below so the self-updater's tag-pinned reinstall keeps
# operating on this one shared install rather than spawning a /root copy.
UV_TOOL_DIR="/opt/axol/uv/tools"
UV_PYTHON_INSTALL_DIR="/opt/axol/uv/python"
export UV_TOOL_DIR UV_PYTHON_INSTALL_DIR

say() { printf "\033[1;36m[axol]\033[0m %s\n" "$*"; }
die() { printf "\033[1;31m[axol]\033[0m %s\n" "$*" >&2; exit 1; }

[ "$(uname -s)" = "Linux" ] || die "This installer only supports Linux."

# The systemd service runs as root, so install everything as root. When piped
# into a non-root shell, re-fetch ourselves and re-run elevated.
if [ "$(id -u)" -ne 0 ]; then
    say "Root privileges are required (systemd service); re-running with sudo..."
    exec sudo bash -c "curl -fsS ${INSTALL_URL} | bash"
fi

command -v systemctl >/dev/null 2>&1 || die "systemd is required."
command -v curl >/dev/null 2>&1 || die "curl is required."
# git is needed both to resolve the latest release tag and by uv itself to
# install from the GitHub URL.
command -v git >/dev/null 2>&1 || die "git is required."

# -- uv -----------------------------------------------------------------------

if ! command -v uv >/dev/null 2>&1; then
    say "Installing uv to ${BIN_DIR}..."
    curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR="${BIN_DIR}" sh
else
    say "uv already installed ($(command -v uv))."
fi
UV="$(command -v uv || echo "${BIN_DIR}/uv")"

# -- axol ---------------------------------------------------------------------

# Installs track releases, not main: resolve the newest release tag (vX.Y.Z)
# and install that exact version from PyPI (the release workflow publishes
# every release there). The `axol serve` self-updater performs the same
# version-pinned reinstall when a newer release appears (serve/update.py) —
# keep the requirement here and there in sync.
say "Resolving the latest release..."
LATEST_TAG="$(git ls-remote --tags "${REPO_URL}" "refs/tags/v*" \
    | awk -F/ '{print $NF}' | sed 's/\^{}$//' \
    | grep -E '^v[0-9]+(\.[0-9]+)*$' | sort -u -V | tail -n1)"
[ -n "${LATEST_TAG}" ] || die "No release tags (vX.Y.Z) found on ${REPO_URL}."
VERSION="${LATEST_TAG#v}"
say "Latest release: ${LATEST_TAG}."

say "Installing axol ${VERSION} from PyPI (extras: ${EXTRAS}) — this can take a few minutes..."
mkdir -p "${UV_TOOL_DIR}" "${UV_PYTHON_INSTALL_DIR}"
UV_TOOL_BIN_DIR="${BIN_DIR}" "${UV}" tool install \
    --python "${PYTHON_VERSION}" \
    --force \
    "almond-axol[${EXTRAS}]==${VERSION}"

# uv creates the tool venv + managed Python with the root umask; make sure every
# user can read/traverse them (dirs + already-executable files only, via 'X') so
# the `axol` wrapper on PATH works for the operator's own login, not just root.
chmod -R a+rX "${UV_TOOL_DIR}" "${UV_PYTHON_INSTALL_DIR}"

command -v axol >/dev/null 2>&1 || die "axol did not end up on PATH (expected ${BIN_DIR}/axol)."
say "axol installed at $(command -v axol)."

# -- System provisioning (pyzed + GStreamer camera stack + adb) ---------------

# `uv tool install` handles the Python package + PyPI deps, but a few pieces it
# can't: the pyzed bindings (not on PyPI), PyGObject (the in-process appsink
# reader for the GPU-resident zedxonesrc/zedsrc camera path, built against the
# system gobject-introspection), and our patched zedxonesrc/zedsrc plugins
# (which stamp each frame at the true sensor-exposure instant so collected
# images line up with joint samples). `axol provision` installs all of them
# idempotently -- the same command the `axol serve` self-updater runs after an
# upgrade, so there is exactly one provisioning path. Each step self-gates
# (a no-op without the ZED SDK / apt / NVENC).
# Also installs adb + the Oculus udev rule (Quest-over-USB controller transport)
# and, on a factory-flashed ZED Box Duo, replaces the known-bad factory GMSL
# camera driver with the pinned release (`axol zed.driver`); the new driver
# only loads at boot, so provision prints a REBOOT REQUIRED notice we surface
# in the summary below rather than rebooting the box mid-install.
say "Provisioning system dependencies (pyzed + GStreamer camera stack + adb)..."
PROVISION_LOG="$(mktemp)"
axol provision 2>&1 | tee "${PROVISION_LOG}" \
    || say "WARNING: provisioning failed; cameras may be unavailable or fall back to the slower ZED SDK path. Re-run 'axol provision'."
REBOOT_REQUIRED=0
if grep -q "REBOOT REQUIRED" "${PROVISION_LOG}"; then REBOOT_REQUIRED=1; fi
rm -f "${PROVISION_LOG}"

# -- Quest-over-USB operator access -------------------------------------------

# `axol provision` writes the Oculus udev rule, which hands the headset to the
# `dialout` group -- the one operators already have for CAN/serial access, so
# interactive `adb` works without a re-login. Make sure the operator is in it
# (provision does this too, but be explicit here where the account is known):
# SUDO_USER is whoever launched the installer (we re-exec under sudo above);
# fall back to the first /home owner when run directly as root.
QUEST_USER="${SUDO_USER:-}"
if [ -z "${QUEST_USER}" ] || [ "${QUEST_USER}" = "root" ]; then
    QUEST_USER="$(ls -ld /home/*/ 2>/dev/null | awk 'NR==1 {print $3}')"
fi
if [ -n "${QUEST_USER}" ] && [ "${QUEST_USER}" != "root" ]; then
    usermod -aG dialout "${QUEST_USER}" 2>/dev/null \
        && say "Added ${QUEST_USER} to 'dialout' for Quest-over-USB adb access." \
        || true
fi

# -- Jetson clocks ------------------------------------------------------------

# Pin the NVENC/VIC engine clocks (encode latency) and the CPU governor (IK
# loop rate). These reset on reboot, so the systemd unit below re-pins them at
# every boot via ExecStartPre; do it once now too so it takes effect without a
# reboot. No-op off a Jetson.
say "Pinning Jetson clocks for the real-time loops..."
axol jetson.setup || say "WARNING: jetson.setup failed (non-Jetson, or needs root)."

# -- Dataset storage location -------------------------------------------------

# The service runs as root, so without help LeRobot would write collected
# datasets to /root/.cache/huggingface/lerobot. Point HF_LEROBOT_HOME at the
# installing user's home instead so datasets live in their account. SUDO_USER
# is set because we re-exec under sudo above; fall back to the owner of the
# first /home/* directory when the installer is run directly as root.
DATASET_USER="${SUDO_USER:-}"
if [ -z "${DATASET_USER}" ] || [ "${DATASET_USER}" = "root" ]; then
    DATASET_USER="$(ls -ld /home/*/ 2>/dev/null | awk 'NR==1 {print $3}')"
fi
LEROBOT_HOME_ENV=""
if [ -n "${DATASET_USER}" ]; then
    DATASET_HOME="$(getent passwd "${DATASET_USER}" | cut -d: -f6)"
    if [ -n "${DATASET_HOME}" ] && [ -d "${DATASET_HOME}" ]; then
        LEROBOT_HOME="${DATASET_HOME}/.cache/huggingface/lerobot"
        say "Datasets will be stored under ${LEROBOT_HOME} (user: ${DATASET_USER})."
        mkdir -p "${LEROBOT_HOME}"
        # Own the cache chain so the user can browse/manage datasets (episode
        # files written by the root service stay root-owned but readable).
        chown "${DATASET_USER}:" \
            "${DATASET_HOME}/.cache" \
            "${DATASET_HOME}/.cache/huggingface" \
            "${LEROBOT_HOME}" 2>/dev/null || true
        LEROBOT_HOME_ENV="Environment=HF_LEROBOT_HOME=${LEROBOT_HOME}"
    fi
fi
if [ -z "${LEROBOT_HOME_ENV}" ]; then
    say "WARNING: could not resolve a non-root user home; datasets default to root's cache."
fi

# -- systemd service ----------------------------------------------------------

say "Writing ${SERVICE_FILE}..."
cat > "${SERVICE_FILE}" <<EOF
[Unit]
Description=Axol control panel (axol serve)
After=network-online.target
Wants=network-online.target

[Service]
# Re-pin the Jetson clocks at every boot (they reset on reboot). The leading
# '-' makes a failure non-fatal, so serve still starts on non-Jetson hosts.
ExecStartPre=-${BIN_DIR}/axol jetson.setup
ExecStart=${BIN_DIR}/axol serve
Restart=always
RestartSec=3
# A stop during a recording session must outlive the serve layer's dataset
# teardown: its stop watchdog gives the recorder subprocess up to 200s to
# finalize the dataset (encode the last episode, flush parquet, write meta).
# systemd's default 90s would SIGKILL the whole cgroup mid-finalize — episodes
# still buffered in the recorder are lost, and resuming the dataset numbers
# the next episode past them, permanently corrupting its episode indexing.
TimeoutStopSec=240
# Include the sbin dirs: CAN bring-up shells out to modinfo/modprobe/ip/etc.
Environment=PATH=${BIN_DIR}:/usr/sbin:/usr/bin:/sbin:/bin
Environment=PYTHONUNBUFFERED=1
# Store collected datasets under the installing user's home (see above) rather
# than root's cache. Empty when no non-root user home was resolved.
${LEROBOT_HOME_ENV}
# Keep the self-updater's tag-pinned reinstall on the same shared install the
# installer wrote, rather than letting uv fall back to its /root defaults.
# UV_TOOL_BIN_DIR must match the install-time bin dir: without it, the first
# self-update relocates the \`axol\` wrapper to uv's default (/root/.local/bin),
# deleting ${BIN_DIR}/axol and breaking this unit's ExecStart with 203/EXEC.
Environment=UV_TOOL_DIR=${UV_TOOL_DIR}
Environment=UV_PYTHON_INSTALL_DIR=${UV_PYTHON_INSTALL_DIR}
Environment=UV_TOOL_BIN_DIR=${BIN_DIR}

[Install]
WantedBy=multi-user.target
EOF

say "Enabling and starting the ${SERVICE_NAME} service..."
systemctl daemon-reload
systemctl enable "${SERVICE_NAME}" >/dev/null 2>&1
systemctl restart "${SERVICE_NAME}"

# -- done ---------------------------------------------------------------------

LAN_IP="$(hostname -I 2>/dev/null | awk "{print \$1}")"
say ""
say "Done! axol serve is running and will start automatically at boot."
say ""
say "  Control panel : https://axol.almond.bot"
if [ -n "${LAN_IP}" ]; then
    say "  This machine  : https://${LAN_IP}:8001 (open once to accept the certificate)"
fi
say ""
say "  Service logs  : journalctl -u ${SERVICE_NAME} -f"
say "  Re-run this installer anytime to upgrade."
if [ "${REBOOT_REQUIRED}" = "1" ]; then
    say ""
    say "REBOOT REQUIRED: the ZED Box camera driver was upgraded and only loads"
    say "at boot. Reboot when convenient:  sudo reboot"
fi
