#!/usr/bin/env bash
# AIMem + OpenShell bring-up: host-side mount, policy apply, optional sandbox launch.
#
# Usage:
#   ./bring-up.sh <bucket> <endpoint-url> [-- <agent command>...]
#
# Env alternatives:
#   BUCKET=<bucket> ENDPOINT_URL=<endpoint-url> ./bring-up.sh [-- <agent command>...]
#
# Optional env:
#   AIMEM_BIN       AIMem binary to run (default: aimem)
#   OPENSHELL_BIN   OpenShell CLI to run (default: openshell)
#   STAGING_DIR     Host path for AIMem's local store (default: /var/cache/aimem-staging)
#   MOUNTPOINT      Host path exposed inside the sandbox (default: /sandbox/aimem-workspace)
#   POLICY_NAME     OpenShell policy name to register (default: aimem-workspace)
#   POLICY_FILE     Pre-rendered OpenShell policy file (default: generated to match MOUNTPOINT)

set -euo pipefail

usage() {
    cat <<'EOF'
AIMem + OpenShell bring-up: host-side mount, policy apply, optional sandbox launch.

Usage:
  ./bring-up.sh <bucket> <endpoint-url> [-- <agent command>...]

Env alternatives:
  BUCKET=<bucket> ENDPOINT_URL=<endpoint-url> ./bring-up.sh [-- <agent command>...]

Optional env:
  AIMEM_BIN       AIMem binary to run (default: aimem)
  OPENSHELL_BIN   OpenShell CLI to run (default: openshell)
  STAGING_DIR     Host path for AIMem's local store (default: /var/cache/aimem-staging)
  MOUNTPOINT      Host path exposed inside the sandbox (default: /sandbox/aimem-workspace)
  POLICY_NAME     OpenShell policy name to register (default: aimem-workspace)
  POLICY_FILE     Pre-rendered OpenShell policy file (default: generated to match MOUNTPOINT)
EOF
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
    usage
    exit 0
fi

BUCKET="${BUCKET:-}"
ENDPOINT_URL="${ENDPOINT_URL:-${AWS_ENDPOINT_URL_S3:-${AWS_ENDPOINT_URL:-}}}"

if [[ $# -gt 0 && "${1}" != "--" && -z "${BUCKET}" ]]; then
    BUCKET="$1"
    shift
fi

if [[ $# -gt 0 && "${1}" != "--" && -z "${ENDPOINT_URL}" ]]; then
    ENDPOINT_URL="$1"
    shift
fi

AGENT_CMD=()
if [[ "${1:-}" == "--" ]]; then
    shift
    if [[ $# -eq 0 ]]; then
        echo "agent command is empty after --" >&2
        usage >&2
        exit 2
    fi
    AGENT_CMD=("$@")
    shift "$#"
fi

if [[ $# -gt 0 ]]; then
    echo "unexpected argument: $1" >&2
    usage >&2
    exit 2
fi

: "${BUCKET:?bucket is required; pass it as the first arg or set BUCKET}"
: "${ENDPOINT_URL:?endpoint URL is required; pass it as the second arg or set ENDPOINT_URL}"

AIMEM_BIN="${AIMEM_BIN:-aimem}"
OPENSHELL_BIN="${OPENSHELL_BIN:-openshell}"
STAGING_DIR="${STAGING_DIR:-/var/cache/aimem-staging}"
MOUNTPOINT="${MOUNTPOINT:-/sandbox/aimem-workspace}"
POLICY_NAME="${POLICY_NAME:-aimem-workspace}"
POLICY_FILE="${POLICY_FILE:-}"

require_command() {
    local cmd="$1"
    if [[ "${cmd}" == */* ]]; then
        [[ -x "${cmd}" ]]
    else
        command -v "${cmd}" >/dev/null 2>&1
    fi
}

missing=()
for cmd in "${AIMEM_BIN}" "${OPENSHELL_BIN}"; do
    if ! require_command "${cmd}"; then
        missing+=("${cmd}")
    fi
done
if [[ ${#missing[@]} -gt 0 ]]; then
    echo "missing required command(s): ${missing[*]}" >&2
    exit 127
fi

GENERATED_POLICY=""
cleanup() {
    if [[ -n "${GENERATED_POLICY}" ]]; then
        rm -f "${GENERATED_POLICY}"
    fi
}
trap cleanup EXIT

stat_pair() {
    stat -c '%d %i' "$1" 2>/dev/null || stat -f '%d %i' "$1" 2>/dev/null
}

is_mountpoint() {
    local path parent path_dev path_ino parent_dev parent_ino

    path="$(cd "$1" 2>/dev/null && pwd -P)" || return 1
    if [[ -r /proc/self/mountinfo ]]; then
        awk -v target="${path}" '
            {
                mount = $5
                gsub(/\\040/, " ", mount)
                gsub(/\\011/, "\t", mount)
                gsub(/\\012/, "\n", mount)
                gsub(/\\134/, "\\", mount)
                if (mount == target) {
                    found = 1
                }
            }
            END { exit found ? 0 : 1 }
        ' /proc/self/mountinfo
        return $?
    fi

    parent="$(cd "${path}/.." 2>/dev/null && pwd -P)" || return 1
    read -r path_dev path_ino < <(stat_pair "${path}") || return 1
    read -r parent_dev parent_ino < <(stat_pair "${parent}") || return 1

    [[ "${path_dev}" != "${parent_dev}" || "${path_ino}" == "${parent_ino}" ]]
}

mkdir -p "${STAGING_DIR}" "${MOUNTPOINT}"

if [[ -z "${POLICY_FILE}" ]]; then
    GENERATED_POLICY="$(mktemp "${TMPDIR:-/tmp}/aimem-openshell-policy.XXXXXX")"
    POLICY_FILE="${GENERATED_POLICY}"
    cat >"${POLICY_FILE}" <<EOF
version: 1

filesystem_policy:
  include_workdir: true

  read_only:
    - /usr
    - /lib
    - /lib64
    - /bin
    - /sbin
    - /etc
    - /proc

  read_write:
    - /tmp
    - "${MOUNTPOINT}"

landlock:
  compatibility: hard_requirement

process:
  run_as_user: sandbox
  run_as_group: sandbox
EOF
fi

AIMEM_PID=""
if is_mountpoint "${MOUNTPOINT}"; then
    echo "==> AIMem mountpoint already live at ${MOUNTPOINT}; leaving it in place" >&2
    echo "    WARNING: cannot verify this mount matches requested bucket '${BUCKET}'." >&2
else
    echo "==> Mounting AIMem bucket '${BUCKET}' at ${MOUNTPOINT}"
    "${AIMEM_BIN}" "${BUCKET}" "${MOUNTPOINT}" \
        --endpoint-url "${ENDPOINT_URL}" \
        --local "${STAGING_DIR}" \
        --foreground &
    AIMEM_PID=$!

    # Wait for the mount to become readable.
    for _ in $(seq 1 30); do
        if is_mountpoint "${MOUNTPOINT}"; then
            echo "==> AIMem live at ${MOUNTPOINT}"
            break
        fi
        if ! kill -0 "${AIMEM_PID}" 2>/dev/null; then
            echo "aimem exited before the mount became ready" >&2
            exit 1
        fi
        sleep 1
    done
    if ! is_mountpoint "${MOUNTPOINT}"; then
        echo "mount never came up" >&2
        kill "${AIMEM_PID}" 2>/dev/null || true
        wait "${AIMEM_PID}" 2>/dev/null || true
        exit 1
    fi
fi

echo "==> Registering OpenShell policy '${POLICY_NAME}'"
"${OPENSHELL_BIN}" policy set "${POLICY_NAME}" \
    --policy "${POLICY_FILE}" \
    --wait

if [[ ${#AGENT_CMD[@]} -gt 0 ]]; then
    echo "==> Launching sandbox with policy '${POLICY_NAME}'"
    if [[ -n "${AIMEM_PID}" ]]; then
        echo "Background AIMem pid: ${AIMEM_PID}. Send SIGTERM to unmount." >&2
    fi
    "${OPENSHELL_BIN}" sandbox run --policy "${POLICY_NAME}" -- "${AGENT_CMD[@]}"
    exit $?
fi

cat <<EOF
==> Done.

Launch a sandbox with the mounted workspace:

    ${OPENSHELL_BIN} sandbox run --policy ${POLICY_NAME} -- /path/to/agent

The agent will see ${MOUNTPOINT} as its writable workspace. AIMem installs
its skill on mount into \$HOME/.agents/skills/aimem and \$HOME/.claude/skills/aimem
(the Claude target relocates to \$CLAUDE_CONFIG_DIR/skills/aimem when
CLAUDE_CONFIG_DIR is set), so the agent picks it up there.

EOF

if [[ -n "${AIMEM_PID}" ]]; then
    echo "Background AIMem pid: ${AIMEM_PID}. Send SIGTERM to unmount."
else
    echo "AIMem was already mounted at ${MOUNTPOINT}; no new process was started."
fi