#!/usr/bin/env bash
# Starts an ephemeral agent sandbox on a buildbox and collects its tool gaps when
# it exits. No agent compute, no container, and no browser ever runs here.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
SCRIPT_PATH="$(readlink -f "${BASH_SOURCE[0]}")"
REGISTRY="$SCRIPT_DIR/../lib/buildbox-registry.mjs"
# shellcheck source=../lib/agent-credentials.sh
. "$SCRIPT_DIR/../lib/agent-credentials.sh"

GAP_STATE="${SANDBOX_GAP_STATE:-$HOME/.local/state/overdeck/sandbox/toolgap.jsonl}"
HOST=""
ID=""
RUNTIME=""
GIT_COMMON=""
GIT_DIR=""
WORKSPACE=""
PROC_STATUS="/proc/self/status"
CGROUP_PATH="/proc/self/cgroup"
PREFLIGHT_ROOT="/run/user"
PRINT_PLAN=0
PASSTHROUGH=()
CMD=()
LAUNCH_ARGS=()
[ "$#" -gt 0 ] && LAUNCH_ARGS=("$@")
RUN_ID="$(date +%s%N)-$$-$RANDOM"

RC_USAGE=2
RC_HOST=3
RC_CREDENTIAL=10
RC_PREFLIGHT=11
RC_SYSTEMD_RUN=12
RC_RECURSION=13
RC_GIT_MOUNT=14

CRED_STAGE_SCRIPT=""
CRED_REMOTE_DIR=""
CRED_HOME_TOKEN=""
CRED_REMOTE_READY=0
SYSTEMD_UNIT=""

SSH_OPTS=(-F "$HOME/.ssh/config" -o BatchMode=yes)
RSYNC_SSH="ssh -F $HOME/.ssh/config -o BatchMode=yes"

usage() {
  cat >&2 <<'USAGE'
usage: agent-sandbox [--host <name>] --id <slug> [--runtime <codex|claude|cursor-agent>]
                     [--git-common <abs path>] [--git-dir <abs path>] [--workspace <abs path>]
                     [--memory 12g] [--pids 512] [--cpus 4]
                     [--proc-status <path>] [--cgroup-path <path>] [--preflight-root <dir>] [--print-plan]
                     [-- <cmd...>]

Without a command the sandbox opens an interactive shell. The workspace lives at
/sandbox/workspaces/<slug> inside the container and survives across runs; the
container itself never does.

--runtime stages that runtime's host credential into the container read-only.
--git-common/--git-dir are mounted read-write at the same absolute path inside
the container; git breaks without them in a dispatched linked worktree.
USAGE
  exit "${1:-$RC_USAGE}"
}

die() { # rc, message...
  local rc="$1"; shift
  printf 'agent-sandbox: %s\n' "$*" >&2
  exit "$rc"
}

add_mount() { PASSTHROUGH+=(--mount "$1:$2:$3"); }

# The credential mount source is node-side absolute, but resolving the node's home
# would cost an ssh round trip; the marker is replaced by a $HOME the node's own
# shell expands. It is unguessable per run, so no argument can forge it.
remote_command_string() {
  local arg out=""
  for arg in "${REMOTE_CMD[@]}"; do out+=" $(printf '%q' "$arg")"; done
  if [ -n "$CRED_HOME_TOKEN" ]; then
    [[ "$out" == *"$CRED_HOME_TOKEN"* ]] ||
      die "$RC_CREDENTIAL" "the credential mount lost its remote-home marker; refusing to mount an unresolved path"
    out="${out//"$CRED_HOME_TOKEN"/\"\$HOME\"}"
  fi
  printf '%s' "$out"
}
set_workspace() { PASSTHROUGH+=(--workspace "$1"); }

require_mount_path() { # flag, path
  case "$2" in
    /*) ;;
    *) die "$RC_GIT_MOUNT" "--$1 must be an absolute path: $2" ;;
  esac
  case "$2" in
    *:*) die "$RC_GIT_MOUNT" "--$1 must not contain ':' — it would corrupt the mount triple: $2" ;;
  esac
}

# An unreadable status file cannot prove the re-exec is unnecessary, so it is
# treated as NoNewPrivs: 1 and the contained path is taken.
no_new_privs() {
  local line value
  [ -r "$PROC_STATUS" ] || { printf '1'; return; }
  while IFS= read -r line; do
    case "$line" in
      NoNewPrivs:*)
        value="${line#NoNewPrivs:}"
        printf '%s' "${value//[[:space:]]/}"
        return
        ;;
    esac
  done <"$PROC_STATUS"
  printf '1'
}

# Recursion is refused on cgroup membership as well as on the environment marker,
# so losing the marker cannot re-arm an unbounded re-exec loop.
inside_reexec_unit() {
  [ -r "$CGROUP_PATH" ] || return 1
  grep -q '/agent-sandbox-[^/]*\.service' "$CGROUP_PATH"
}

user_manager_preflight() {
  local uid bus
  uid="$(id -u)"
  bus="$PREFLIGHT_ROOT/$uid/bus"

  [ "$(loginctl show-user "$USER" -p Linger --value 2>/dev/null)" = yes ] ||
    die "$RC_PREFLIGHT" "user manager preflight: lingering is not enabled for $USER — run: loginctl enable-linger $USER"
  [ "${XDG_RUNTIME_DIR:-}" = "$PREFLIGHT_ROOT/$uid" ] ||
    die "$RC_PREFLIGHT" "user manager preflight: XDG_RUNTIME_DIR is '${XDG_RUNTIME_DIR:-unset}', expected $PREFLIGHT_ROOT/$uid — run: loginctl enable-linger $USER"
  [ -S "$bus" ] ||
    die "$RC_PREFLIGHT" "user manager preflight: no user bus socket at $bus — run: loginctl enable-linger $USER"
}

# shellcheck disable=SC2329  # invoked from the EXIT trap
cleanup() {
  local rc=$?
  if [ "${AGENT_SANDBOX_REEXEC:-}" = 1 ] && [ -n "${AGENT_SANDBOX_REEXEC_STATUS_FILE:-}" ]; then
    printf '%s\n' "$rc" >"$AGENT_SANDBOX_REEXEC_STATUS_FILE" 2>/dev/null || true
  fi
  if [ "$CRED_REMOTE_READY" = 1 ]; then
    # shellcheck disable=SC2029  # the path is quoted here on purpose
    ssh "${SSH_OPTS[@]}" "$HOST" "rm -rf -- $(printf '%q' "$CRED_REMOTE_DIR")" >/dev/null 2>&1 || true
  fi
  return "$rc"
}
trap cleanup EXIT
trap 'forward_signal; exit 130' INT
trap 'forward_signal; exit 143' TERM

# shellcheck disable=SC2329  # invoked from the INT/TERM traps
forward_signal() {
  [ -n "$SYSTEMD_UNIT" ] || return 0
  systemctl --user stop "$SYSTEMD_UNIT" >/dev/null 2>&1 || true
}

# systemd-run --wait --pipe reports the payload's status, but it reports its own
# failures the same way; the payload records its status so the two are distinct.
reexec_through_systemd_run() {
  local status_file rc recorded
  status_file="$(mktemp -p "$XDG_RUNTIME_DIR" "agent-sandbox-status.XXXXXX")"
  SYSTEMD_UNIT="agent-sandbox-$RUN_ID"

  set +e
  systemd-run --user --wait --pipe --collect --expand-environment=no \
    --unit="$SYSTEMD_UNIT" \
    --setenv=PATH="$PATH" \
    --setenv=HOME="$HOME" \
    --setenv=AGENT_SANDBOX_REEXEC=1 \
    --setenv=AGENT_SANDBOX_REEXEC_STATUS_FILE="$status_file" \
    -- "$SCRIPT_PATH" ${LAUNCH_ARGS[@]+"${LAUNCH_ARGS[@]}"}
  rc=$?
  set -e
  SYSTEMD_UNIT=""

  recorded="$(cat "$status_file" 2>/dev/null || true)"
  rm -f "$status_file"
  if [[ ! "$recorded" =~ ^[0-9]+$ ]]; then
    die "$RC_SYSTEMD_RUN" "systemd-run did not run the launcher (exit $rc); the contained process never started"
  fi
  exit "$recorded"
}

while [ $# -gt 0 ]; do
  case "$1" in
    --host)           HOST="${2-}"; shift 2 ;;
    --id)             ID="${2-}"; PASSTHROUGH+=(--id "${2-}"); shift 2 ;;
    --runtime)        RUNTIME="${2-}"; shift 2 ;;
    --git-common)     GIT_COMMON="${2-}"; shift 2 ;;
    --git-dir)        GIT_DIR="${2-}"; shift 2 ;;
    --workspace)      WORKSPACE="${2-}"; shift 2 ;;
    --memory)         PASSTHROUGH+=(--memory "${2-}"); shift 2 ;;
    --pids)           PASSTHROUGH+=(--pids "${2-}"); shift 2 ;;
    --cpus)           PASSTHROUGH+=(--cpus "${2-}"); shift 2 ;;
    --proc-status)    PROC_STATUS="${2-}"; shift 2 ;;
    --cgroup-path)    CGROUP_PATH="${2-}"; shift 2 ;;
    --preflight-root) PREFLIGHT_ROOT="${2-}"; shift 2 ;;
    --print-plan)     PRINT_PLAN=1; shift ;;
    --)               shift; CMD=("$@"); break ;;
    -h|--help)        usage 0 ;;
    *) echo "agent-sandbox: unknown arg: $1" >&2; usage ;;
  esac
done

[ -n "$ID" ] || { echo "agent-sandbox: --id is required" >&2; usage; }
[ -r "$REGISTRY" ] || die "$RC_USAGE" "buildbox registry resolver missing at $REGISTRY"

if [ -n "$WORKSPACE" ]; then
  require_mount_path workspace "$WORKSPACE"
  set_workspace "$WORKSPACE"
fi
if [ -n "$GIT_COMMON" ]; then
  require_mount_path git-common "$GIT_COMMON"
  add_mount "$GIT_COMMON" "$GIT_COMMON" rw
fi
if [ -n "$GIT_DIR" ]; then
  require_mount_path git-dir "$GIT_DIR"
  add_mount "$GIT_DIR" "$GIT_DIR" rw
fi

if [ "$(no_new_privs)" = 1 ]; then
  if [ "${AGENT_SANDBOX_REEXEC:-}" = 1 ] || inside_reexec_unit; then
    die "$RC_RECURSION" "already re-executed through systemd-run yet NoNewPrivs is still 1; refusing to recurse"
  fi
  user_manager_preflight
  reexec_through_systemd_run
fi

if [ -z "$HOST" ]; then
  HOST="$(node -e '
    import("'"$REGISTRY"'").then(({ loadRegistry, resolveHosts }) => {
      const registry = loadRegistry();
      const byName = new Map(registry.hosts.map((h) => [h.name, h]));
      const candidates = resolveHosts(registry, { order: "build" })
        .filter((name) => byName.get(name).roles.includes("agent-seat"));
      if (candidates.length === 0) throw new Error("no enabled registry host carries the agent-seat role");
      process.stdout.write(byName.get(candidates[0]).ssh_alias);
    }).catch((err) => { process.stderr.write(`agent-sandbox: ${err.message}\n`); process.exit(3); });
  ')"
fi
[ -n "$HOST" ] || exit "$RC_HOST"

if [ -n "$RUNTIME" ]; then
  CRED_LINE="$(agent_cred_resolve "$RUNTIME")" || exit "$RC_CREDENTIAL"
  IFS=$'\t' read -r CRED_RUNTIME CRED_REALPATH CRED_CONTAINER_PATH <<<"$CRED_LINE"
  printf 'agent-sandbox: %s credential identity: %s\n' "$CRED_RUNTIME" "$CRED_REALPATH" >&2

  CRED_REMOTE_DIR=".local/state/overdeck-sandbox/creds/$RUN_ID"
  CRED_HOME_TOKEN="@@agent-sandbox-remote-home-$RUN_ID@@"

  add_mount "$CRED_HOME_TOKEN/$CRED_REMOTE_DIR/$CRED_RUNTIME/${CRED_CONTAINER_PATH##*/}" \
    "$(agent_cred_staging_path "$CRED_RUNTIME")" ro
  CRED_STAGE_SCRIPT="$(agent_cred_stage_script "$CRED_RUNTIME")"
fi

REMOTE_CMD=(".local/share/overdeck-sandbox/bin/sandbox-run" "${PASSTHROUGH[@]}")
if [ "${#CMD[@]}" -gt 0 ]; then
  REMOTE_CMD+=(--)
  if [ -n "$CRED_STAGE_SCRIPT" ]; then
    # "$@" rather than exec: the stage script's cleanup trap only fires if this
    # shell outlives the agent.
    REMOTE_CMD+=(/bin/sh -c "$CRED_STAGE_SCRIPT"$'\n''"$@"' sh)
  fi
  REMOTE_CMD+=("${CMD[@]}")
fi

if [ "$PRINT_PLAN" = 1 ]; then
  printf 'agent-sandbox plan: host=%s\n' "$HOST"
  printf 'agent-sandbox plan: remote-cmd%s\n' "$(remote_command_string)"
  exit 0
fi

if [ -t 0 ] && [ -t 1 ]; then
  SSH_OPTS+=(-t)
  REMOTE_CMD=("${REMOTE_CMD[0]}" --tty "${REMOTE_CMD[@]:1}")
fi

if [ -n "$RUNTIME" ]; then
  # shellcheck disable=SC2029  # the path is quoted here on purpose
  ssh "${SSH_OPTS[@]}" "$HOST" \
    "umask 077 && mkdir -p -- $(printf '%q' "$CRED_REMOTE_DIR/$CRED_RUNTIME")" ||
    die "$RC_CREDENTIAL" "cannot create the credential staging directory on $HOST"
  CRED_REMOTE_READY=1
  rsync -q -e "$RSYNC_SSH" --chmod=F600 \
    "$CRED_REALPATH" "$HOST:$CRED_REMOTE_DIR/$CRED_RUNTIME/${CRED_CONTAINER_PATH##*/}" ||
    die "$RC_CREDENTIAL" "cannot stage the $CRED_RUNTIME credential on $HOST"
fi

set +e
# shellcheck disable=SC2029  # the remote-home marker is meant to expand on the node
ssh "${SSH_OPTS[@]}" "$HOST" "$(remote_command_string)"
RC=$?
set -e

STATE_DIR="$(dirname "$GAP_STATE")"
mkdir -p "$STATE_DIR"
rsync -q -e "$RSYNC_SSH" "$HOST:.local/share/overdeck-sandbox/agent-image-tag" \
  "$STATE_DIR/agent-image-tag" 2>/dev/null || true
PER_HOST_DIR="$STATE_DIR/hosts"
mkdir -p "$PER_HOST_DIR"
if rsync -q -e "$RSYNC_SSH" "$HOST:sandbox/toolgap/gaps.jsonl" "$PER_HOST_DIR/$HOST.jsonl" 2>/dev/null; then
  cat "$PER_HOST_DIR"/*.jsonl >"$GAP_STATE.new"
  mv "$GAP_STATE.new" "$GAP_STATE"
fi

exit "$RC"