#!/usr/bin/env bash
# Host-side AIMem mount for OpenSandbox.
#
# OpenSandbox sandboxes expose no /dev/fuse and no CAP_SYS_ADMIN, so aimem
# cannot mount inside a sandbox. This mounts aimem on the HOST; main.py then
# binds the mountpoint into the sandbox as a `host` volume at /workspace —
# the same model OpenSandbox uses for its ossfs volumes.
#
# Usage:
#   aimem-host-mount.sh [mount] <sandbox-id>
#   aimem-host-mount.sh cleanup <sandbox-id>
#
#   Prints the host mountpoint on the last stdout line once it is live and
#   leaves `aimem --foreground` running in the background.
#
# Env: AIMEM_BUCKET, AIMEM_ENDPOINT (required);
#      AIMEM_ACCESS_KEY / AIMEM_SECRET_KEY (+ optional AIMEM_SESSION_TOKEN),
#      read by aimem from the environment;
#      AIMEM_LOCAL (default /tmp/aimem-<sandbox-id>-local),
#      AIMEM_AGENT (default auto), AIMEM_MOUNT_ROOT (default /mnt/aimem),
#      AIMEM_VERSION (default latest).
set -euo pipefail

usage() {
    echo "usage: $0 [mount] <sandbox-id> | cleanup <sandbox-id>" >&2
}

MOUNT_ROOT="${AIMEM_MOUNT_ROOT:-/mnt/aimem}"
DL_BASE="${AIMEM_DL_BASE:-https://dl.min.io/aistor/aimem/release}"
VERSION="${AIMEM_VERSION:-latest}"

# sudo isn't guaranteed on an OpenSandbox host — the process managing
# sandboxes there often already runs as root, and minimal host images may
# not ship sudo at all. Only shell out to it when actually needed.
SUDO=""
[ "$(id -u)" -ne 0 ] && SUDO="sudo"

cleanup_mount() {
    # Two separate `local` statements: `local a=X b=$a` expands both
    # right-hand sides before either assignment takes effect, so `$a` in b's
    # value is unbound under `set -u`, not the freshly-declared `a`.
    local id="$1"
    local mountpoint="${MOUNT_ROOT}/${id}"

    if mountpoint -q "$mountpoint"; then
        if command -v fusermount3 >/dev/null 2>&1 && fusermount3 -u "$mountpoint"; then
            :
        elif command -v fusermount >/dev/null 2>&1 && fusermount -u "$mountpoint"; then
            :
        elif ${SUDO} umount "$mountpoint"; then
            :
        else
            echo "failed to unmount ${mountpoint}" >&2
            return 1
        fi
        echo "aimem unmounted at ${mountpoint}" >&2
    fi

    rmdir "$mountpoint" 2>/dev/null || true
}

if [ "$#" -eq 1 ]; then
    CMD="mount"
    ID="$1"
elif [ "$#" -eq 2 ]; then
    CMD="$1"
    ID="$2"
else
    usage
    exit 2
fi

MOUNTPOINT="${MOUNT_ROOT}/${ID}"

case "$CMD" in
    mount)
        ;;
    cleanup)
        cleanup_mount "$ID"
        exit 0
        ;;
    *)
        usage
        exit 2
        ;;
esac

: "${AIMEM_BUCKET:?AIMEM_BUCKET required}"
: "${AIMEM_ENDPOINT:?AIMEM_ENDPOINT required}"

case "$(uname -m)" in
    x86_64|amd64) AIMEM_ARCH="amd64" ;;
    aarch64|arm64) AIMEM_ARCH="arm64" ;;
    *)
        echo "unsupported host architecture: $(uname -m)" >&2
        exit 1
        ;;
esac

# Install a released binary from the public CDN if absent, verifying it
# against its published .sha256sum.
install_bin() {
    local name="$1" dest="$2" tmp marker
    marker="${dest}.version"
    # Not just `[ -x "$dest" ]`: that alone never re-checks $VERSION, so a
    # host that installed an older/staged binary once would silently keep
    # running it forever regardless of a later AIMEM_VERSION change.
    if [ -x "$dest" ] && [ "$(cat "$marker" 2>/dev/null)" = "$VERSION" ]; then
        return 0
    fi
    tmp="$(mktemp)"
    # latest is served under the bare file name; a pinned tag lives under
    # archive/, the only path that still resolves once a newer release ships.
    local suffix="" dir=""
    [ "$VERSION" = latest ] || { suffix=".${VERSION}"; dir="archive/"; }
    curl -fsSL \
        "${DL_BASE}/linux-${AIMEM_ARCH}/${dir}${name}${suffix}" \
        -o "$tmp"
    curl -fsSL \
        "${DL_BASE}/linux-${AIMEM_ARCH}/${dir}${name}${suffix}.sha256sum" \
        -o "${tmp}.sum"
    echo "$(cut -d' ' -f1 "${tmp}.sum")  ${tmp}" | sha256sum -c -
    ${SUDO} install -m 0755 "$tmp" "$dest"
    echo "$VERSION" | ${SUDO} tee "$marker" >/dev/null
    rm -f "$tmp" "${tmp}.sum"
}

install_bin aimem            /usr/local/bin/aimem
install_bin aimem-fusermount /usr/bin/aimem-fusermount

# The mount runs as this (host) user; aimem forwards the /dev/fuse fd to the
# setuid-root helper for the privileged mount(2), so it must be root:root 4755.
${SUDO} chown root:root /usr/bin/aimem-fusermount
${SUDO} chmod 4755      /usr/bin/aimem-fusermount

aimem --version

# The mountpoint must sit under a prefix in sandbox.toml [storage]
# allowed_host_paths (see sandbox.toml.example) or OpenSandbox rejects the
# host-volume bind at validation time.
mkdir -p "$MOUNTPOINT"

LOG="/tmp/aimem-${ID}.log"
# allow_root is a config-file-only tunable: the container runtime that binds
# this mountpoint into a sandbox (dockerd, containerd, ...) runs as root, and a
# FUSE mount is otherwise only visible to the mounting user — without it, the
# container runtime can't traverse the mount to set up the bind and fails
# opaquely (e.g. Docker's "error while creating mount source path: file exists"
# when it can stat but not fully access the mount).
CONFIG="$(mktemp -t aimem-opensandbox-XXXXXX.yaml)"
printf 'version: "1.0"\nallow_root: true\n' >"$CONFIG"
aimem "${AIMEM_BUCKET}" "${MOUNTPOINT}" \
    --endpoint-url "${AIMEM_ENDPOINT}" \
    --local "${AIMEM_LOCAL:-/tmp/aimem-${ID}-local}" \
    --agent "${AIMEM_AGENT:-auto}" \
    --config "${CONFIG}" \
    --foreground >"$LOG" 2>&1 &
MOUNT_PID=$!

for _ in $(seq 1 30); do
    if mountpoint -q "$MOUNTPOINT"; then
        echo "aimem mounted at ${MOUNTPOINT} (pid ${MOUNT_PID})" >&2
        echo "$MOUNTPOINT"   # last stdout line = mountpoint, consumed by main.py
        exit 0
    fi
    # Surface the log instead of hanging if aimem died during startup.
    if ! kill -0 "$MOUNT_PID" 2>/dev/null; then
        echo "aimem exited early:" >&2
        cat "$LOG" >&2
        exit 1
    fi
    sleep 1
done

echo "timed out waiting for ${MOUNTPOINT}; mount log:" >&2
cat "$LOG" >&2
kill "$MOUNT_PID" 2>/dev/null || true
exit 1