#!/usr/bin/env bash
# Installs the agent sandbox onto a buildbox: host tooling, the forced-command
# e2e channel, and the pinned images. Touches only the target user's home —
# never /etc, never a systemd unit, never network or sshd configuration.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
MODULE_DIR="$(cd "$SCRIPT_DIR/../../../sandbox" && pwd)"
REGISTRY="$SCRIPT_DIR/../lib/buildbox-registry.mjs"
SHARED_PAIR_LIB="$SCRIPT_DIR/../lib/e2e-pair.sh"
HASH_LIB="$MODULE_DIR/host/lib/image-context.sh"
AUTHORIZED_KEYS_MARKER="overdeck-sandbox-e2e"
SANDBOX_ROLE="agent-sandbox"
HOSTS=()
ALL=0
CHECK=0
FORCE_BUILD=()

usage() {
  cat >&2 <<'USAGE'
usage: sandbox-provision [--host <name>]... [--all] [--check] [--rebuild]

Resolves the target from the buildbox registry (order "build", role "agent-seat")
unless --host names one explicitly; --host may be repeated. Idempotent.

--all    every reachable registry host carrying the "agent-sandbox" role.
--check  report only: compares each box's built image tag against the tag this
         checkout's image context produces. Mutates nothing.

exit 0: every probed host carries the expected image
exit 1: a reachable host is unprovisioned or serving an older image (--check)
exit 4: nothing drifted, but a host could not be reached (--check)
USAGE
  exit "${1:-2}"
}

while [ $# -gt 0 ]; do
  case "$1" in
    --host)    HOSTS+=("${2-}"); shift 2 ;;
    --all)     ALL=1; shift ;;
    --check)   CHECK=1; shift ;;
    --rebuild) FORCE_BUILD=(--force); shift ;;
    -h|--help) usage 0 ;;
    *) echo "sandbox-provision: unknown arg: $1" >&2; usage ;;
  esac
done

[ -r "$REGISTRY" ] || { echo "sandbox-provision: buildbox registry resolver missing at $REGISTRY" >&2; exit 2; }
[ -r "$SHARED_PAIR_LIB" ] || { echo "sandbox-provision: e2e pair orchestrator missing at $SHARED_PAIR_LIB" >&2; exit 2; }
[ -r "$HASH_LIB" ] || { echo "sandbox-provision: context digest helper missing at $HASH_LIB" >&2; exit 2; }
# shellcheck source=../../../sandbox/host/lib/image-context.sh
. "$HASH_LIB"

[ "$ALL" = 0 ] || [ "${#HOSTS[@]}" -eq 0 ] || { echo "sandbox-provision: --all and --host are exclusive" >&2; usage; }

# `node` on PATH is the cpu-guard shim, which may offload the run to a buildbox and
# execute it from a mirrored path there — where this checkout's resolver does not exist,
# and where a file written by the script would not be this machine's. Resolving which
# hosts to provision must happen on this machine.
if [[ -x /usr/bin/node ]]; then NODE_BIN=/usr/bin/node; else NODE_BIN="$(command -v node || true)"; fi
[ -n "$NODE_BIN" ] || { echo "sandbox-provision: no node interpreter found" >&2; exit 2; }

# Prints one ssh alias per line for the registry hosts carrying $2, in $1 order.
# The result travels through a file, not stdout, so that anything else writing to this
# process's stdout can never be mistaken for a host name.
registry_aliases() { # order role [first-only]
  local out="$STAGE_ALIASES"
  "$NODE_BIN" -e '
    import("'"$REGISTRY"'").then(({ loadRegistry, resolveHosts }) => {
      const registry = loadRegistry();
      const byName = new Map(registry.hosts.map((h) => [h.name, h]));
      const names = resolveHosts(registry, { order: process.argv[1] })
        .filter((name) => byName.get(name).roles.includes(process.argv[2]));
      if (names.length === 0) throw new Error(`no enabled registry host carries the ${process.argv[2]} role`);
      const picked = process.argv[3] === "first" ? names.slice(0, 1) : names;
      require("node:fs").writeFileSync(process.argv[4], picked.map((n) => byName.get(n).ssh_alias).join("\n") + "\n");
    }).catch((err) => { process.stderr.write(`sandbox-provision: ${err.message}\n`); process.exit(3); });
  ' "$1" "$2" "${3-}" "$out" >/dev/null || return 3
  [ -s "$out" ] || return 3
  cat "$out"
}

STAGE="$(mktemp -d -t sandbox-provision.XXXXXX)"
trap 'rm -rf "$STAGE"' EXIT
STAGE_ALIASES="$STAGE/registry-aliases"

# Process substitution discards the resolver's exit status, so the resolution is
# captured and its status checked: an unresolvable registry must stop the run, never
# leave it probing an empty or malformed host list.
resolved=""
if [ "$ALL" = 1 ]; then
  resolved="$(registry_aliases build "$SANDBOX_ROLE")" || exit 3
elif [ "${#HOSTS[@]}" -eq 0 ]; then
  resolved="$(registry_aliases build agent-seat first)" || exit 3
fi
if [ -n "$resolved" ]; then
  readarray -t HOSTS <<<"$resolved"
fi
[ "${#HOSTS[@]}" -gt 0 ] || exit 3
for host in "${HOSTS[@]}"; do
  [[ "$host" =~ ^[A-Za-z0-9._-]+$ ]] \
    || { echo "sandbox-provision: registry resolution produced a non-host line: $host" >&2; exit 3; }
done

mkdir -p "$STAGE/bin" "$STAGE/lib" "$STAGE/image-context"
cp "$MODULE_DIR"/host/bin/* "$STAGE/bin/"
cp "$MODULE_DIR"/host/lib/*.mjs "$MODULE_DIR"/host/lib/*.sh "$MODULE_DIR"/host/lib/*.py \
  "$MODULE_DIR"/host/lib/*.json "$STAGE/lib/"
cp "$MODULE_DIR"/lib/sandbox-payload.mjs "$STAGE/lib/"
cp -r "$MODULE_DIR"/image/. "$STAGE/image-context/"
# The base context may already carry generated files. Unlock it before overlays replace them.
chmod -R u+w "$STAGE/image-context"
mkdir -p "$STAGE/image-context/lib"
cp "$MODULE_DIR"/lib/sandbox-payload.mjs "$STAGE/image-context/lib/"
cp "$SHARED_PAIR_LIB" "$STAGE/image-context/lib/e2e-pair.sh"
# Release trees are intentionally immutable. Recursive copy preserves their read-only modes,
# but this staging tree must accept generated overlays, produce mode-stable images, and remain
# removable by the caller. Normalize after every source file has entered the stage.
chmod -R u+w "$STAGE/bin" "$STAGE/lib" "$STAGE/image-context"
chmod 0755 "$STAGE"/bin/* "$STAGE"/image-context/bin/*

CONTEXT_HASH="$(sandbox_image_context_hash "$STAGE/image-context")"
EXPECTED_TAG="localhost/overdeck-agent-sandbox:$CONTEXT_HASH"
EXPECTED_E2E_TAG="localhost/overdeck-agent-sandbox-e2e:$CONTEXT_HASH"

ROOT='$HOME/.local/share/overdeck-sandbox'
# This probe runs inside deploy-local.sh while it holds the deploy lock, so every wait it
# can incur is bounded: ConnectTimeout covers only TCP and auth, the keepalives cover a
# connection that dies mid-command, and the remote timeout covers a podman image store
# already locked by a build on that box.
ssh_to() {
  ssh -F "$HOME/.ssh/config" -o BatchMode=yes -o ConnectTimeout=10 \
    -o ServerAliveInterval=10 -o ServerAliveCountMax=3 "$@"
}

LOG_DIR="$STAGE/logs"
mkdir -p "$LOG_DIR"

# Runs one independent operation per host, then replays each host's output only after
# every child has finished. This keeps a slow or failed box from hiding the result of
# the others, while preserving readable, non-interleaved reports.
run_host_phase() { # phase worker hosts...
  local phase=$1 worker=$2 host log i failed=0
  local -a phase_hosts=() pids=() results=()
  LAST_PHASE_SUCCESS_HOSTS=()
  LAST_PHASE_FAILED_HOSTS=()
  shift 2
  for host in "$@"; do
    log="$LOG_DIR/$phase.$host.log"
    ( "$worker" "$host" ) >"$log" 2>&1 &
    phase_hosts+=("$host")
    pids+=("$!")
  done
  for i in "${!pids[@]}"; do
    if wait "${pids[$i]}"; then
      results+=(0)
      LAST_PHASE_SUCCESS_HOSTS+=("${phase_hosts[$i]}")
    else
      results+=(1)
      LAST_PHASE_FAILED_HOSTS+=("${phase_hosts[$i]}")
      failed=1
    fi
  done
  for i in "${!phase_hosts[@]}"; do
    cat "$LOG_DIR/$phase.${phase_hosts[$i]}.log"
    if [ "${results[$i]}" != 0 ] && [ "$CHECK" != 1 ]; then
      printf 'sandbox-provision: %s FAILED during %s\n' "${phase_hosts[$i]}" "$phase" >&2
    fi
  done
  return "$failed"
}

if [ "$CHECK" = 1 ]; then
  echo "sandbox-provision: expected image $EXPECTED_TAG"
  # One round trip per host reports both images: sandbox-run serves the agent image and
  # the e2e pair serves its sibling, and they are built from the same context, so a box
  # carrying only one of them is drifted too.
  probe='R="$HOME/.local/share/overdeck-sandbox"
    for f in agent-image-tag e2e-image-tag; do
      t="$(cat "$R/$f" 2>/dev/null || true)"
      if [ -n "$t" ] && timeout 15 podman image exists "$t" 2>/dev/null; then e=present; else e=absent; fi
      printf "image-record %s %s %s\n" "$f" "${t:-none}" "$e"
    done'
  check_host() {
    local host=$1 report host_ok seen marker kind have present want
    if ! report="$(ssh_to "$host" "$probe" 2>/dev/null)"; then
      printf '%s unreachable\n' "$host"
      return 4
    fi
    host_ok=1
    # A remote shell sources the box's rc files, so it may print lines of its own. Only
    # the marked records are read, and both must arrive — an absent record is drift, not
    # silence.
    seen=0
    while read -r marker kind have present; do
      [ "$marker" = image-record ] || continue
      seen=$((seen + 1))
      want="$EXPECTED_TAG"
      [ "$kind" = agent-image-tag ] || want="$EXPECTED_E2E_TAG"
      if [ "$have" = none ]; then
        printf '%s unprovisioned — no %s built\n' "$host" "$kind"; host_ok=0
      elif [ "$have" != "$want" ]; then
        printf '%s drifted %s have=%s want=%s\n' "$host" "$kind" "$have" "$want"; host_ok=0
      elif [ "$present" != present ]; then
        printf '%s drifted %s=%s recorded but the image is gone\n' "$host" "$kind" "$have"; host_ok=0
      fi
    done <<<"$report"
    if [ "$seen" != 2 ]; then
      printf '%s probe returned %s of 2 image records\n' "$host" "$seen"; host_ok=0
    fi
    if [ "$host_ok" = 1 ]; then
      printf '%s converged %s\n' "$host" "$EXPECTED_TAG"
    else
      :
    fi
    [ "$host_ok" = 1 ]
  }
  if ! run_host_phase probe check_host "${HOSTS[@]}"; then
    # Preserve the check contract: any drift is 1, while unreachable hosts alone are 4.
    check_failed=0
    check_unreachable=0
    for host in "${HOSTS[@]}"; do
      if grep -qxF "$host unreachable" "$LOG_DIR/probe.$host.log"; then
        check_unreachable=1
      else
        check_failed=1
      fi
    done
    [ "$check_failed" = 0 ] || exit 1
    [ "$check_unreachable" = 0 ] || exit 4
  fi
  exit 0
fi

probe_host() {
  local HOST=$1
  local -a SSH=(ssh -F "$HOME/.ssh/config" -o BatchMode=yes "$HOST")
  "${SSH[@]}" "mkdir -p $ROOT/bin $ROOT/lib $ROOT/image-context $ROOT/secrets \
    \$HOME/sandbox/workspaces \$HOME/sandbox/home \$HOME/sandbox/store/pnpm \$HOME/sandbox/toolgap \
    && chmod 0700 $ROOT/secrets"
}

sync_host() {
  local HOST=$1
  # One connection carries all three trees; --delete prunes within each copied dir
  # exactly as the previous per-dir transfers did.
  rsync -a --delete -e "ssh -F $HOME/.ssh/config -o BatchMode=yes" \
    "$STAGE/bin" "$STAGE/lib" "$STAGE/image-context" "$HOST:.local/share/overdeck-sandbox/"
}

build_host() {
  local HOST=$1 SSH_PORT=$2 BUILT BUILT_E2E
  local -a SSH=(ssh -F "$HOME/.ssh/config" -o BatchMode=yes "$HOST")
  "${SSH[@]}" "SANDBOX_BUILD_FORCE='${FORCE_BUILD[*]-}' SANDBOX_SSH_PORT='$SSH_PORT' bash -s" <<'REMOTE'
set -euo pipefail
ROOT="$HOME/.local/share/overdeck-sandbox"
KEY="$ROOT/secrets/e2e_key"
MARKER="overdeck-sandbox-e2e"
AK="$HOME/.ssh/authorized_keys"

[ -f "$KEY" ] || ssh-keygen -q -t ed25519 -N '' -C "$MARKER" -f "$KEY"
chmod 0600 "$KEY"

printf '[host.containers.internal]:%s %s\n' "$SANDBOX_SSH_PORT" \
  "$(cut -d' ' -f1,2 /etc/ssh/ssh_host_ed25519_key.pub)" >"$ROOT/secrets/e2e_known_hosts"
chmod 0644 "$ROOT/secrets/e2e_known_hosts"
printf '%s\n' "$SANDBOX_SSH_PORT" >"$ROOT/e2e-port"

tr -d '\n' </etc/machine-id >"$ROOT/host-machine-id"

mkdir -p "$HOME/.ssh"
touch "$AK"
chmod 0600 "$AK"
LINE="restrict,command=\"$ROOT/bin/sandbox-e2e-receiver\" $(cat "$KEY.pub")"
TMP="$(mktemp)"
grep -v -F "$MARKER" "$AK" >"$TMP" || true
printf '%s\n' "$LINE" >>"$TMP"
chmod 0600 "$TMP"
mv "$TMP" "$AK"

if "$ROOT/bin/sandbox-image-build" --pull; then
  :
else
  pull_rc=$?
  if [ "$pull_rc" -eq 4 ]; then
    echo "sandbox-provision: registry pull miss; falling back to local build" >&2
    "$ROOT/bin/sandbox-image-build" --local-only ${SANDBOX_BUILD_FORCE:-}
  else
    echo "sandbox-provision: registry pull was refused; local fallback is unsafe" >&2
    exit "$pull_rc"
  fi
fi
REMOTE

  BUILT="$("${SSH[@]}" "cat $ROOT/agent-image-tag")"
  [ "$BUILT" = "$EXPECTED_TAG" ] \
    || { echo "sandbox-provision: $HOST built $BUILT but this checkout expects $EXPECTED_TAG" >&2; return 1; }
  BUILT_E2E="$("${SSH[@]}" "cat $ROOT/e2e-image-tag")"
  [ "$BUILT_E2E" = "$EXPECTED_E2E_TAG" ] \
    || { echo "sandbox-provision: $HOST built $BUILT_E2E but this checkout expects $EXPECTED_E2E_TAG" >&2; return 1; }
  echo "sandbox-provision: $HOST provisioned ($BUILT)"
}

# Each host runs its whole probe -> rsync -> build pipeline independently: a fast box
# never waits at a phase barrier for the slowest one, and the failing step names
# itself in the host's own log.
provision_host() {
  local HOST=$1 SSH_PORT
  # The sandbox reaches the receiver through podman's host mapping, so it needs the
  # port sshd actually answers on. This client config is the value already proven to work.
  SSH_PORT="$(ssh -F "$HOME/.ssh/config" -G "$HOST" | awk '/^port /{print $2; exit}')"
  [[ "$SSH_PORT" =~ ^[0-9]+$ ]] || { echo "sandbox-provision: could not resolve the ssh port for $HOST" >&2; return 2; }
  echo "sandbox-provision: target $HOST (sshd port $SSH_PORT)"
  probe_host "$HOST" || { echo "sandbox-provision: $HOST failed during probe" >&2; return 1; }
  sync_host "$HOST" || { echo "sandbox-provision: $HOST failed during rsync" >&2; return 1; }
  build_host "$HOST" "$SSH_PORT" || { echo "sandbox-provision: $HOST failed during build" >&2; return 1; }
}

run_host_phase provision provision_host "${HOSTS[@]}"
exit $?