#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# Run the Rust e2e smoke test against an openshell-gateway running the
# standalone VM compute driver (`openshell-driver-vm`).
#
# Architecture:
#   * `openshell-sandbox` runs inside a NIC-less guest.
#   * The host `openshell-supervisor` connects over virtio-vsock and owns
#     gateway registration, policy evaluation, DNS, and external networking.
#   * Readiness is authoritative on the gateway: a sandbox's phase
#     flips to `Ready` the moment `ConnectSupervisor` registers, and
#     back to `Provisioning` when the session drops. The VM driver
#     only reports `Error` conditions for dead launcher processes.
#
# Usage:
#   mise run e2e:vm
#
# What the script does:
#   1. When no prebuilt VM driver is supplied, ensures the VM runtime
#      (libkrun) and bundled sandbox/supervisor binaries are staged.
#   2. Builds `openshell-gateway`, `openshell-driver-vm`, the native host
#      `openshell-supervisor` host control, and the `openshell` CLI with the
#      embedded runtime as needed. When CI supplies OPENSHELL_GATEWAY_BIN,
#      OPENSHELL_VM_DRIVER_BIN, OPENSHELL_VM_SUPERVISOR_BIN, or OPENSHELL_BIN,
#      the matching prebuilt binary is reused instead of rebuilt.
#   3. On macOS, codesigns the VM driver (libkrun needs the
#      `com.apple.security.hypervisor` entitlement).
#   4. Writes a per-run gateway config with `[openshell.drivers.vm]`
#      settings, starts the gateway with `--config <run-state>/gateway.toml`
#      on a random free port, waits for an authenticated gateway status
#      request to succeed, then runs the selected Rust e2e tests.
#   5. Tears the gateway down and (on failure) preserves the gateway
#      log and every VM serial console log for post-mortem.
#
# Prerequisites (handled automatically by this script for local VM-driver builds):
#   - `mise run vm:setup`      — downloads / builds the libkrun runtime.
#   - `mise run vm:supervisor` — builds the bundled sandbox supervisor.

set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
source "${ROOT}/e2e/support/gateway-common.sh"
# shellcheck source=e2e/support/conformance.sh
source "${ROOT}/e2e/support/conformance.sh"

COMPRESSED_DIR="${ROOT}/target/vm-runtime-compressed"
GATEWAY_BIN="${OPENSHELL_GATEWAY_BIN:-${ROOT}/target/debug/openshell-gateway}"
DRIVER_BIN="${OPENSHELL_VM_DRIVER_BIN:-${ROOT}/target/debug/openshell-driver-vm}"
CLI_BIN="${OPENSHELL_BIN:-${ROOT}/target/debug/openshell}"
E2E_TEST_OVERRIDE="${OPENSHELL_E2E_VM_TEST:-}"
# Unset or empty selects Rust defaults; e2e-conformance-only skips Rust tests.
E2E_FEATURES="${OPENSHELL_E2E_VM_FEATURES:-e2e-vm}"
SANDBOX_IMAGE="${OPENSHELL_SANDBOX_IMAGE:-nvcr.io/nvidia/base/ubuntu:24.04}"

# The VM driver places `compute-driver.sock` under `[openshell.drivers.vm].state_dir`.
# AF_UNIX SUN_LEN is 104 bytes on macOS (108 on Linux), so paths anchored
# in the workspace's `target/` blow the limit on typical developer
# machines — e.g. a ~100-char `~/.superset/worktrees/.../target/...`
# prefix plus the `compute-driver.sock` leaf leaves no room. macOS'
# per-user `$TMPDIR` (`/var/folders/xx/.../T/`) can be 50+ chars too,
# so root state under `/tmp` unconditionally to keep UDS paths short.
STATE_DIR_ROOT="/tmp"

# Smoke test timeouts. First boot extracts the embedded libkrun runtime
# (~60-90MB of zstd per architecture) and prepares an ext4 root disk from the
# configured image. The guest then starts the sandbox supervisor directly; a cold
# microVM is typically ready within ~15s after image preparation.
GATEWAY_READY_TIMEOUT=60
SANDBOX_PROVISION_TIMEOUT=180

# ── Build prerequisites ──────────────────────────────────────────────

if [ -n "${RUSTC_WRAPPER:-}" ] && [ "${OPENSHELL_E2E_VM_ALLOW_RUSTC_WRAPPER:-0}" != "1" ]; then
  echo "==> Building without RUSTC_WRAPPER=${RUSTC_WRAPPER} (set OPENSHELL_E2E_VM_ALLOW_RUSTC_WRAPPER=1 to keep it)"
  unset RUSTC_WRAPPER
fi

if [ -z "${OPENSHELL_VM_DRIVER_BIN:-}" ]; then
  mkdir -p "${COMPRESSED_DIR}"

  if ! find "${COMPRESSED_DIR}" -maxdepth 1 -name 'libkrun*.zst' | grep -q .; then
    echo "==> Preparing embedded VM runtime (mise run vm:setup)"
    mise run vm:setup
  fi

  # Always rebuild the guest bundle so an e2e run cannot silently exercise a
  # stale boundary binary after supervisor or isolation-interface changes.
  echo "==> Building bundled VM supervisor (mise run vm:supervisor)"
  mise run vm:supervisor

  export OPENSHELL_VM_RUNTIME_COMPRESSED_DIR="${OPENSHELL_VM_RUNTIME_COMPRESSED_DIR:-${COMPRESSED_DIR}}"
else
  echo "==> Skipping embedded VM runtime prep; prebuilt VM driver was supplied"
fi

build_packages=()
if [ -z "${OPENSHELL_GATEWAY_BIN:-}" ]; then
  if [ "${OPENSHELL_E2E_EXTERNAL_COMPUTE_DRIVER:-0}" = "1" ]; then
    echo "==> Building driver-free openshell-gateway"
    cargo build \
      -p openshell-gateway --bin openshell-gateway \
      --no-default-features --features telemetry
  else
    build_packages+=(-p openshell-gateway)
  fi
else
  echo "==> Using prebuilt openshell-gateway at ${GATEWAY_BIN}"
fi
if [ -z "${OPENSHELL_VM_DRIVER_BIN:-}" ]; then
  build_packages+=(-p openshell-driver-vm)
else
  echo "==> Using prebuilt openshell-driver-vm at ${DRIVER_BIN}"
fi
if [ -z "${OPENSHELL_VM_SUPERVISOR_BIN:-}" ]; then
  # The VM driver prefers a native sibling `openshell-supervisor`. Build it
  # explicitly so a stale target/debug binary cannot disagree with the
  # freshly embedded guest sandbox protocol.
  build_packages+=(-p openshell-supervisor)
else
  echo "==> Using prebuilt VM host supervisor at ${OPENSHELL_VM_SUPERVISOR_BIN}"
fi
if [ -z "${OPENSHELL_BIN:-}" ]; then
  build_packages+=(-p openshell-cli)
else
  echo "==> Using prebuilt openshell CLI at ${CLI_BIN}"
fi

if [ "${#build_packages[@]}" -gt 0 ]; then
  echo "==> Building ${build_packages[*]}"
  cargo build "${build_packages[@]}"
fi

for pair in \
  "openshell-gateway:${GATEWAY_BIN}" \
  "openshell-driver-vm:${DRIVER_BIN}" \
  "openshell CLI:${CLI_BIN}"; do
  label="${pair%%:*}"
  path="${pair#*:}"
  if [ ! -x "${path}" ]; then
    echo "ERROR: expected ${label} binary at ${path}" >&2
    exit 1
  fi
done
export OPENSHELL_BIN="${CLI_BIN}"
DRIVER_DIR="$(dirname "${DRIVER_BIN}")"

if [ "$(uname -s)" = "Darwin" ]; then
  echo "==> Codesigning openshell-driver-vm (Hypervisor entitlement)"
  codesign \
    --entitlements "${ROOT}/crates/openshell-driver-vm/entitlements.plist" \
    --force \
    -s - \
    "${DRIVER_BIN}"
fi

# ── Pick a random free host port for the gateway ─────────────────────

HOST_PORT="$(python3 -c 'import socket
s = socket.socket()
s.bind(("", 0))
print(s.getsockname()[1])
s.close()')"

# Per-run state dir so concurrent e2e runs don't collide on the UDS or
# sandbox state. The VM driver creates `<state_dir>/compute-driver.sock`
# and `<state_dir>/sandboxes/<id>/overlay.ext4` under here. Keep the
# basename short — see the SUN_LEN comment above.
RUN_STATE_DIR="${STATE_DIR_ROOT}/os-vm-e2e-${HOST_PORT}-$$"
mkdir -p "${RUN_STATE_DIR}"
export XDG_CONFIG_HOME="${RUN_STATE_DIR}/config"
export XDG_DATA_HOME="${RUN_STATE_DIR}/data"

GATEWAY_LOG="$(mktemp /tmp/openshell-gateway-e2e.XXXXXX)"
GATEWAY_PID_FILE="${RUN_STATE_DIR}/gateway.pid"
GATEWAY_ARGS_FILE="${RUN_STATE_DIR}/gateway.args"
GATEWAY_CONFIG="${RUN_STATE_DIR}/gateway.toml"
GATEWAY_DB="${RUN_STATE_DIR}/gateway.db"
JWT_DIR="${RUN_STATE_DIR}/jwt"
PKI_DIR="${RUN_STATE_DIR}/pki"
GATEWAY_NAME="openshell-e2e-vm-${HOST_PORT}"
DRIVER_PID=""
DRIVER_LOG="${RUN_STATE_DIR}/vm-driver.log"
DRIVER_SOCKET="${RUN_STATE_DIR}/compute-driver.sock"

# ── Cleanup (trap) ───────────────────────────────────────────────────

cleanup() {
  local exit_code=$?

  local gateway_pid="${GATEWAY_PID:-}"
  if [ -f "${GATEWAY_PID_FILE:-}" ]; then
    gateway_pid="$(cat "${GATEWAY_PID_FILE}" 2>/dev/null || true)"
  fi

  if [ -n "${gateway_pid}" ] && kill -0 "${gateway_pid}" 2>/dev/null; then
    echo "Stopping openshell-gateway (pid ${gateway_pid})..."
    # SIGTERM first; gateway drops ManagedDriverProcess which SIGKILLs
    # the driver and removes the UDS. Wait briefly, then force-kill.
    kill -TERM "${gateway_pid}" 2>/dev/null || true
    for _ in 1 2 3 4 5 6 7 8 9 10; do
      kill -0 "${gateway_pid}" 2>/dev/null || break
      sleep 0.5
    done
    kill -KILL "${gateway_pid}" 2>/dev/null || true
    wait "${gateway_pid}" 2>/dev/null || true
  fi
  e2e_stop_process "${DRIVER_PID}" "external VM compute driver"

  # On failure, keep the VM console log for debugging. We deliberately
  # print it instead of leaving it on disk because the state dir gets
  # wiped on success.
  if [ "${exit_code}" -ne 0 ]; then
    echo "=== gateway log (preserved for debugging) ==="
    cat "${GATEWAY_LOG}" 2>/dev/null || true
    echo "=== end gateway log ==="
    if [ -f "${DRIVER_LOG}" ]; then
      echo "=== external VM compute driver log ==="
      cat "${DRIVER_LOG}" 2>/dev/null || true
      echo "=== end external VM compute driver log ==="
    fi

    local console
    while IFS= read -r -d '' console; do
      echo "=== VM console log: ${console} ==="
      cat "${console}" 2>/dev/null || true
      echo "=== end VM console log ==="
    done < <(find "${RUN_STATE_DIR}/sandboxes" -name 'rootfs-console.log' -print0 2>/dev/null)
  fi

  rm -f "${GATEWAY_LOG}" 2>/dev/null || true
  # Only wipe the per-run state dir on success. On failure, leave it for
  # post-mortem (serial console logs and root disk images).
  if [ "${exit_code}" -eq 0 ]; then
    rm -rf "${RUN_STATE_DIR}" 2>/dev/null || true
  else
    echo "NOTE: preserving ${RUN_STATE_DIR} for debugging"
  fi
}
trap cleanup EXIT

# ── Launch the gateway + VM driver ───────────────────────────────────

echo "==> Starting openshell-gateway on 127.0.0.1:${HOST_PORT} (state: ${RUN_STATE_DIR})"

# Pin `driver_dir` to the workspace `target/debug/` so we always pick up
# the driver we just cargo-built. Without this, the gateway's
# `resolve_compute_driver_bin` fallback prefers
# `~/.local/libexec/openshell/openshell-driver-vm` when present,
# which silently shadows development builds — a subtle source of
# stale-binary bugs in e2e runs.
# `grpc_endpoint` is consumed by the host supervisor. The host alias is
# normalized to loopback while keeping package-managed certificate naming.
e2e_generate_gateway_jwt "${JWT_DIR}"
e2e_generate_pki "${GATEWAY_BIN}" "${PKI_DIR}"

cat >"${GATEWAY_CONFIG}" <<EOF
[openshell]
version = 2

[openshell.gateway]
bind_address = "127.0.0.1:${HOST_PORT}"
compute_driver = "vm"
guest_tls_ca = "${PKI_DIR}/ca.crt"

[openshell.gateway.tls]
cert_path = "${PKI_DIR}/server/tls.crt"
key_path = "${PKI_DIR}/server/tls.key"
client_ca_path = "${PKI_DIR}/ca.crt"

[openshell.gateway.mtls_auth]
enabled = true

[openshell.gateway.gateway_jwt]
signing_key_path = "${JWT_DIR}/signing.pem"
public_key_path = "${JWT_DIR}/public.pem"
kid_path = "${JWT_DIR}/kid"
gateway_id = "${GATEWAY_NAME}"
# Local VM e2e gateways exercise the single-player default: sandbox JWTs
# identify the supervisor and do not expire.

[openshell.drivers.vm]
EOF
# A rootless Linux libkrun refuses file creation on read-write bind mounts by
# a guest UID/GID other than its own, so vm_bind_mount runs the sandbox as the
# gateway user there.
VM_IDENTITY_ARGS=()
if [ "$(uname -s)" = "Linux" ] && [ "$(id -u)" != "0" ]; then
  VM_IDENTITY_ARGS=(--sandbox-uid "$(id -u)" --sandbox-gid "$(id -g)")
fi
if [ "${OPENSHELL_E2E_EXTERNAL_COMPUTE_DRIVER:-0}" = "1" ]; then
  printf 'socket_path = "%s"\n' "${DRIVER_SOCKET}" >>"${GATEWAY_CONFIG}"
else
  cat >>"${GATEWAY_CONFIG}" <<EOF
grpc_endpoint = "https://host.openshell.internal:${HOST_PORT}"
driver_dir = "${DRIVER_DIR}"
state_dir = "${RUN_STATE_DIR}"
enable_bind_mounts = true
EOF
  if [ "${#VM_IDENTITY_ARGS[@]}" -gt 0 ]; then
    printf 'sandbox_uid = %s\nsandbox_gid = %s\n' "$(id -u)" "$(id -g)" >>"${GATEWAY_CONFIG}"
  fi
fi
# Host bind mounts are exercised by vm_bind_mount, mirroring the Docker and
# Podman e2e gateways. An external driver must acknowledge the same policy.
# Keep [openshell.drivers.vm] the last table: vm_corporate_proxy appends bare
# keys to it.
cat >>"${GATEWAY_CONFIG}" <<EOF
allow_driver_config = true
resource_admission = { enabled = false }
EOF

if [ "${OPENSHELL_E2E_EXTERNAL_COMPUTE_DRIVER:-0}" = "1" ]; then
  "${DRIVER_BIN}" \
    --bind-socket "${DRIVER_SOCKET}" \
    --allow-same-uid-peer \
    --grpc-endpoint "https://host.openshell.internal:${HOST_PORT}" \
    --default-image "${SANDBOX_IMAGE}" \
    --state-dir "${RUN_STATE_DIR}" \
    --guest-tls-ca "${PKI_DIR}/ca.crt" \
    --admission-config-json '{"allow_driver_config":true,"resource_admission":{"enabled":false}}' \
    --enable-bind-mounts \
    ${VM_IDENTITY_ARGS[@]+"${VM_IDENTITY_ARGS[@]}"} \
    >"${DRIVER_LOG}" 2>&1 &
  DRIVER_PID=$!
  e2e_wait_for_socket \
    "${DRIVER_SOCKET}" "${DRIVER_PID}" "external VM compute driver" 60
fi

GATEWAY_ARGS=(
  --config "${GATEWAY_CONFIG}"
  --db-url "sqlite:${GATEWAY_DB}?mode=rwc"
)
e2e_write_gateway_args_file "${GATEWAY_ARGS_FILE}" "${GATEWAY_ARGS[@]}"

"${GATEWAY_BIN}" "${GATEWAY_ARGS[@]}" \
  >"${GATEWAY_LOG}" 2>&1 &
GATEWAY_PID=$!
printf '%s\n' "${GATEWAY_PID}" >"${GATEWAY_PID_FILE}"

# Register the gateway before polling so readiness exercises the same mTLS
# client path as the smoke tests.
CLI_GATEWAY_ENDPOINT="https://127.0.0.1:${HOST_PORT}"
e2e_register_mtls_gateway \
  "${XDG_CONFIG_HOME}" \
  "${GATEWAY_NAME}" \
  "${CLI_GATEWAY_ENDPOINT}" \
  "${HOST_PORT}" \
  "${PKI_DIR}"
export OPENSHELL_GATEWAY_ENDPOINT="${CLI_GATEWAY_ENDPOINT}"

# ── Wait for gateway readiness ───────────────────────────────────────
#
# Poll the authenticated gRPC health path instead of coupling readiness to a
# particular gateway log message. The VM driver is spawned eagerly, while
# sandboxes are created on demand.

echo "==> Waiting for gateway readiness (timeout ${GATEWAY_READY_TIMEOUT}s)"
elapsed=0
last_status_output=""
while [ "${elapsed}" -lt "${GATEWAY_READY_TIMEOUT}" ]; do
  if ! kill -0 "${GATEWAY_PID}" 2>/dev/null; then
    echo "ERROR: openshell-gateway exited before becoming ready"
    exit 1
  fi
  if last_status_output="$("${CLI_BIN}" status --output json 2>&1)" &&
    printf '%s\n' "${last_status_output}" |
      grep -Eq '"status"[[:space:]]*:[[:space:]]*"connected"'; then
    echo "==> Gateway ready after ${elapsed}s"
    break
  fi
  sleep 2
  elapsed=$((elapsed + 2))
done

if [ "${elapsed}" -ge "${GATEWAY_READY_TIMEOUT}" ]; then
  echo "ERROR: openshell-gateway did not become ready after ${GATEWAY_READY_TIMEOUT}s"
  echo "=== last openshell status output ==="
  if [ -n "${last_status_output}" ]; then
    printf '%s\n' "${last_status_output}"
  else
    echo "<no output>"
  fi
  echo "=== end openshell status output ==="
  exit 1
fi

# ── Run the smoke test ───────────────────────────────────────────────
#
# The CLI uses the raw endpoint but still resolves matching metadata so it
# can find the mTLS client bundle.

export OPENSHELL_E2E_DRIVER="vm"
export OPENSHELL_E2E_VM_STATE_DIR="${RUN_STATE_DIR}"
e2e_export_gateway_restart_metadata \
  "${GATEWAY_BIN}" \
  "${GATEWAY_ARGS_FILE}" \
  "${GATEWAY_LOG}" \
  "${GATEWAY_PID_FILE}"

# The VM driver creates each sandbox VM from a cached read-only ext4 root disk
# plus a writable overlay disk. The guest's sandbox supervisor then initializes
# policy, netns, Landlock, and sshd. On a cold host this is ~15s after image
# preparation; allow 180s for slower CI runners.
export OPENSHELL_PROVISION_TIMEOUT="${SANDBOX_PROVISION_TIMEOUT}"

e2e_run_openshell_conformance "VM"
if [ "${E2E_FEATURES}" = "e2e-conformance-only" ]; then
  echo "::notice::Skipping Rust e2e tests (e2e-conformance-only)"
  exit 0
fi

# Seed the catalog once for the whole lane. The profiles live in the gateway
# for as long as it runs, and the import is create-only: a second import of
# the same directory fails with "custom provider profile '<id>' already
# exists", so this cannot move inside the per-target helper below.
e2e_import_example_provider_profiles "${CLI_BIN}" "${ROOT}" || exit 1

run_e2e_test() {
  local test_target="$1"
  shift

  echo "==> Running e2e ${test_target} test (features: ${E2E_FEATURES}, endpoint: ${OPENSHELL_GATEWAY_ENDPOINT})"
  cargo test \
    --manifest-path "${ROOT}/e2e/rust/Cargo.toml" \
    --features "${E2E_FEATURES}" \
    --test "${test_target}" \
    "$@" \
    -- --nocapture
  echo "==> ${test_target} test passed."
}

if [ -n "${E2E_TEST_OVERRIDE}" ]; then
  run_e2e_test "${E2E_TEST_OVERRIDE}"
else
  run_e2e_test ephemeral_cleanup
  run_e2e_test host_gateway_alias
  run_e2e_test vm_overlay
  run_e2e_test vm_bind_mount
  run_e2e_test vm_gateway_start
  run_e2e_test vm_corporate_proxy
fi