#!/usr/bin/env bash
# Install aimem into a Daytona linux-vm sandbox at first boot.
#
# Daytona's declarative image builder is unavailable in the us-east-1 region
# and for linux-vm sandboxes, so the snapshot is a stock OS image and aimem is
# installed here at runtime instead of baked into an image layer. aimem links
# only libc/libm/libgcc_s (ibverbs/RDMA is a dlopen-only fast path, unused over
# an HTTP S3 endpoint), so no extra system packages are required for the mount.
#
# usage.py uploads this script plus aimem-mount to /tmp and, on the CI/offline
# path, a locally built binary to /tmp/aimem-staged.
#
# usage.py asks for a root sandbox (`os_user="root"`), but Daytona has been
# observed to start one as the unprivileged `daytona` user instead, so don't
# assume uid 0: escalate with passwordless sudo when it is available, and fail
# with a diagnosable message when neither is.
set -euo pipefail

echo "aimem-setup: running as $(id)"

if [ "$(id -u)" -eq 0 ]; then
    SUDO=""
elif command -v sudo >/dev/null 2>&1 && sudo -n true 2>/dev/null; then
    SUDO="sudo"
    echo "aimem-setup: not root, escalating with passwordless sudo"
else
    echo "aimem-setup: need root or passwordless sudo to install into" \
         "/usr/local/bin, but running as $(id -un) with neither" >&2
    exit 1
fi

# usage.py can only upload into /tmp (the SDK writes as the sandbox user), which
# is world-writable, so don't install straight from there: another process in the
# sandbox could swap a source file or point it at a symlink and have us install
# attacker-controlled code as root. Copy the uploads into a root-owned 0700 dir
# and install from that. This narrows but does not close the window — the copy
# still reads /tmp — so keep the upload set small and never widen it to anything
# the sandbox itself produces.
STAGE="$($SUDO mktemp -d /root/aimem-setup.XXXXXX)"
$SUDO chmod 0700 "$STAGE"
cleanup() { $SUDO rm -rf "$STAGE"; }
trap cleanup EXIT

# Stage one upload. Any problem is fatal and reports its own cause: callers must
# test for presence themselves rather than treating a non-zero return as "absent",
# which would mask a real staging failure.
stage_input() {
    src="/tmp/$1"
    if [ -L "$src" ]; then
        echo "aimem-setup: $src is a symlink; refusing to install from it" >&2
        exit 1
    fi
    if [ ! -f "$src" ]; then
        echo "aimem-setup: $src is missing or not a regular file" >&2
        exit 1
    fi
    if ! $SUDO install -m 0600 "$src" "$STAGE/$1"; then
        echo "aimem-setup: failed to stage $src into $STAGE" >&2
        exit 1
    fi
}

stage_input aimem-mount
$SUDO install -m 0755 "$STAGE/aimem-mount" /usr/local/bin/aimem-mount

AIMEM_VERSION="${AIMEM_VERSION:-latest}"
AIMEM_DL_BASE="${AIMEM_DL_BASE:-https://dl.min.io/aistor/aimem/release}"
case "$(uname -m)" in
    x86_64 | amd64) dl_arch=linux-amd64 ;;
    aarch64 | arm64) dl_arch=linux-arm64 ;;
    *) echo "aimem-setup: unsupported arch $(uname -m)" >&2; exit 1 ;;
esac

DEST=/usr/local/bin/aimem
if [ -e /tmp/aimem-staged ]; then
    stage_input aimem-staged
    $SUDO install -m 0755 "$STAGE/aimem-staged" "$DEST"
else
    if ! command -v curl >/dev/null 2>&1; then
        # sudo resets the environment, so pass DEBIAN_FRONTEND through `env`
        # rather than exporting it.
        $SUDO env DEBIAN_FRONTEND=noninteractive apt-get update
        $SUDO env DEBIAN_FRONTEND=noninteractive apt-get install -y \
            --no-install-recommends ca-certificates curl
    fi
    tmp_bin="$(mktemp)"
    tmp_sum="$(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.
    suffix=""; dir=""
    [ "$AIMEM_VERSION" = latest ] || { suffix=".${AIMEM_VERSION}"; dir="archive/"; }
    curl -fsSL "${AIMEM_DL_BASE}/${dl_arch}/${dir}aimem${suffix}" -o "$tmp_bin"
    curl -fsSL "${AIMEM_DL_BASE}/${dl_arch}/${dir}aimem${suffix}.sha256sum" -o "$tmp_sum"
    echo "$(cut -d' ' -f1 "$tmp_sum")  $tmp_bin" | sha256sum -c -
    $SUDO install -m 0755 "$tmp_bin" "$DEST"
    rm -f "$tmp_bin" "$tmp_sum"
fi

"$DEST" --version
echo "aimem-setup: OK"