#!/bin/bash
# Linux GPU Runner Startup Script
#
# This script runs as root (GCE metadata scripts always run as root).
# The GitHub Actions runner must run as the user who owns the runner
# directory owner), not as root.
#
# Steps:
# 1. Removes any pre-existing runner service from the base image
# 2. Updates the preinstalled GitHub Actions runner if it is stale
# 3. Reads the JIT config from GCP instance metadata
# 4. Starts the GitHub Actions runner as the correct user
# 5. Shuts down the VM when the job completes

set -euo pipefail

RUNNER_VERSION="2.334.0"
RUNNER_SHA256="048024cd2c848eb6f14d5646d56c13a4def2ae7ee3ad12122bee960c56f3d271"

# Defense in depth: wipe stray actions-runner installs that may be baked into
# the base image under user accounts other than the canonical "runner" user.
# A May 2026 SM80Plus outage was caused by /home/<engineer>/actions-runner/
# leaking into the snapshot — the runner-detection loop below picked it up
# and ran the GitHub Actions runner under the engineer's account, registering
# without labels and blocking the queue. Image hygiene is the canonical fix;
# this is a runtime safety net so a single bad snapshot can't repeat that.
for stray in /home/*/actions-runner; do
  # Without `nullglob`, an unmatched glob expands to the literal pattern;
  # the directory check below skips it.
  [ -d "$stray" ] || continue
  case "$stray" in
  /home/runner/actions-runner) continue ;;
  esac
  echo "Wiping stray actions-runner install: $stray"
  # Stop and uninstall the runner service if it is registered. Run the two
  # phases independently — `uninstall` must still be attempted even when
  # `stop` fails (service not running, partially installed, etc.) so we do
  # not leave a dangling systemd unit pointing at the directory we are about
  # to delete.
  if [ -f "$stray/.runner" ] && [ -x "$stray/svc.sh" ]; then
    (cd "$stray" && ./svc.sh stop 2>/dev/null) || true
    (cd "$stray" && ./svc.sh uninstall 2>/dev/null) || true
  fi
  rm -rf "$stray"
done

# Configure host-side core-dump capture for CI containers.
# Containers running under `--user root` without `--privileged` cannot write
# `kernel.core_pattern` themselves (it is silently rejected as read-only),
# so we set it on the host before any job container starts. The CI workflow
# bind-mounts /var/cores into the container with `-v /var/cores:/var/cores`,
# so cores written by a crashing process inside the container land here on
# the host where the workflow's post-test step can read and gdb-extract them.
# Mode 1777 (sticky-writable) so non-root processes inside the container can
# write their own cores even after the workflow drops privileges.
mkdir -p /var/cores
chmod 1777 /var/cores
sysctl -w kernel.core_pattern='/var/cores/core.%e.%p' >/dev/null
sysctl -w kernel.core_uses_pid=1 >/dev/null

# Find the runner directory and its owner
RUNNER_DIR=""
RUNNER_USER=""
for home in /home/*; do
  if [ -d "$home/actions-runner" ]; then
    RUNNER_DIR="$home/actions-runner"
    RUNNER_USER=$(stat -c '%U' "$home/actions-runner")
    break
  fi
done

if [ -z "$RUNNER_DIR" ]; then
  if [ -d "/actions-runner" ]; then
    RUNNER_DIR="/actions-runner"
    RUNNER_USER=$(stat -c '%U' "/actions-runner")
  else
    echo "ERROR: Cannot find actions-runner directory"
    shutdown -h now
    exit 1
  fi
fi

LOG_FILE="${RUNNER_DIR}/startup.log"

log() {
  local msg
  msg="$(date '+%Y-%m-%d %H:%M:%S') - $1"
  echo "$msg"
  echo "$msg" >>"$LOG_FILE"
}

log "=== Linux Runner Startup ==="
log "Runner directory: $RUNNER_DIR"
log "Runner user: $RUNNER_USER"

fail_update_and_shutdown() {
  log "ERROR: $1"
  if [ -n "${runner_archive:-}" ]; then
    rm -f "$runner_archive" || true
  fi
  shutdown -h now
  exit 1
}

# Step 0: Remove any pre-existing runner service from the base image.
log "Removing pre-existing runner service (if any)..."
if systemctl list-units --type=service --all 2>/dev/null | grep -q "actions.runner"; then
  cd "$RUNNER_DIR"
  ./svc.sh stop 2>&1 | while read -r line; do log "  $line"; done || true
  ./svc.sh uninstall 2>&1 | while read -r line; do log "  $line"; done || true
  log "  Service removed."
else
  log "  No existing runner service found."
fi

# Remove old runner config files
for f in .runner .credentials .credentials_rsaparams .runner_migrated; do
  if [ -f "$RUNNER_DIR/$f" ]; then
    rm -f "$RUNNER_DIR/$f"
    log "  Removed $f"
  fi
done

runner_version() {
  if [ -x "$RUNNER_DIR/bin/Runner.Listener" ]; then
    sudo -u "$RUNNER_USER" "$RUNNER_DIR/bin/Runner.Listener" --version 2>/dev/null | head -n 1 | tr -d '\r'
  fi
}

current_runner_version="$(runner_version || true)"
if [ -z "$current_runner_version" ]; then
  current_runner_version="unknown"
fi
log "Current Actions runner version: $current_runner_version"

if [ "$current_runner_version" != "$RUNNER_VERSION" ]; then
  log "Updating Actions runner to v${RUNNER_VERSION}..."
  if ! runner_archive="$(mktemp /tmp/actions-runner.XXXXXX.tar.gz)"; then
    fail_update_and_shutdown "Failed to create temporary Actions runner archive"
  fi
  runner_url="https://github.com/actions/runner/releases/download/v${RUNNER_VERSION}/actions-runner-linux-x64-${RUNNER_VERSION}.tar.gz"

  if ! curl -fsSL --retry 3 --connect-timeout 10 --max-time 120 "$runner_url" -o "$runner_archive"; then
    fail_update_and_shutdown "Failed to download Actions runner v${RUNNER_VERSION}"
  fi

  if ! printf '%s  %s\n' "$RUNNER_SHA256" "$runner_archive" | sha256sum -c - >/dev/null 2>&1; then
    fail_update_and_shutdown "Actions runner v${RUNNER_VERSION} checksum verification failed"
  fi

  if ! chown "$RUNNER_USER":"$RUNNER_USER" "$runner_archive"; then
    fail_update_and_shutdown "Failed to change owner for Actions runner v${RUNNER_VERSION} archive"
  fi

  if ! sudo -u "$RUNNER_USER" tar xzf "$runner_archive" -C "$RUNNER_DIR"; then
    fail_update_and_shutdown "Failed to extract Actions runner v${RUNNER_VERSION}"
  fi
  rm -f "$runner_archive" || true
  runner_archive=""

  updated_runner_version="$(runner_version || true)"
  log "Actions runner version after update: ${updated_runner_version:-unknown}"
  if [ "${updated_runner_version:-}" != "$RUNNER_VERSION" ]; then
    fail_update_and_shutdown "Runner version mismatch after update (expected ${RUNNER_VERSION}, got ${updated_runner_version:-unknown})"
  fi
fi

# Step 0.5: Ensure NVIDIA GPU devices are initialized.
#
# This block only applies to GPU pools (the Linux test runners). CPU-only pools
# (the build and analytics runners) share both this startup script *and* the
# same base image, so the nvidia-smi tool and driver libraries are present even
# when no GPU is attached — tool presence is therefore NOT a reliable signal.
#
# Whether this pool is *supposed* to have a GPU is authoritative, not inferred:
# the scaler stamps an "expect-gpu" metadata key from the pool's --gcp-gpu-type
# config (true for GPU pools, false for CPU-only build/analytics pools). We then
# combine that contract with the actual hardware, detected by a physically-
# attached NVIDIA device on the PCI bus (vendor ID 10de — prefer lspci, fall
# back to the sysfs PCI vendor list if pciutils is absent):
#   - GPU pool + GPU present     -> initialize the GPU (fatal if init fails).
#   - GPU pool + GPU MISSING     -> fatal: the accelerator attach failed; a
#                                   runner here would silently accept GPU jobs.
#   - CPU-only pool              -> skip GPU init and register as a CPU runner.
# Defaulting expect-gpu to "true" keeps existing GPU pools fail-safe even if the
# metadata is somehow absent.
EXPECT_GPU="$(curl -sf --max-time 10 --connect-timeout 5 \
  -H "Metadata-Flavor: Google" \
  "http://metadata.google.internal/computeMetadata/v1/instance/attributes/expect-gpu" 2>/dev/null || echo "true")"
log "GPU expectation for this pool: expect-gpu=${EXPECT_GPU}"

gpu_present=false
if command -v lspci >/dev/null 2>&1; then
  if lspci -d 10de: 2>/dev/null | grep -q .; then
    gpu_present=true
  fi
elif grep -qi '0x10de' /sys/bus/pci/devices/*/vendor 2>/dev/null; then
  gpu_present=true
fi

if [ "$EXPECT_GPU" != "false" ] && [ "$gpu_present" != "true" ]; then
  log "ERROR: This pool expects an NVIDIA GPU but none is attached (accelerator attach failed?). Refusing to register a GPU runner with no device."
  shutdown -h now
  exit 1
fi

if [ "$gpu_present" = "true" ]; then
  log "Initializing NVIDIA GPU..."
  gpu_ready=false
  for attempt in $(seq 1 10); do
    if nvidia-smi >/dev/null 2>&1; then
      log "  GPU initialized successfully."
      gpu_ready=true
      break
    fi
    log "  Attempt ${attempt}/10: nvidia-smi not ready, waiting..."
    sleep 5
  done

  if [ "$gpu_ready" != "true" ]; then
    log "ERROR: GPU initialization failed after 10 attempts"
    shutdown -h now
    exit 1
  fi

  # Create nvidia-modeset device if it doesn't exist (needed for Vulkan)
  if [ ! -e /dev/nvidia-modeset ]; then
    log "  Creating /dev/nvidia-modeset..."
    nvidia-modprobe -m 2>/dev/null || modprobe nvidia-modeset 2>/dev/null || true
  fi

  # Enable GPU persistence mode to prevent NVML state corruption in containers.
  # Without this, NVML can lose track of GPU processes when they exit inside Docker,
  # causing "Failed to initialize NVML: Unknown Error".
  log "  Enabling GPU persistence mode..."
  if pm_out="$(nvidia-smi -pm 1 2>&1)"; then
    log "  GPU persistence mode enabled."
  else
    log "WARNING: Failed to enable GPU persistence mode: ${pm_out}"
  fi

  # Create /dev/char symlinks for all NVIDIA device nodes. Recent runc versions
  # with cgroup v2 require these symlinks to properly inject devices into
  # containers. Without them, containers can intermittently lose GPU access
  # with "Failed to initialize NVML: Unknown Error".
  # See: https://github.com/NVIDIA/nvidia-docker/issues/1730
  log "  Creating /dev/char symlinks..."
  if ctk_out="$(nvidia-ctk system create-dev-char-symlinks --create-all 2>&1)"; then
    log "  /dev/char symlinks created."
  else
    log "WARNING: Failed to create /dev/char symlinks: ${ctk_out}"
  fi

  # Verify GPU devices
  log "  GPU devices:"
  ls -la /dev/nvidia* 2>&1 | while read -r line; do log "    $line"; done || true
  ls -la /dev/dri/* 2>&1 | while read -r line; do log "    $line"; done || true
else
  log "No NVIDIA GPU on the PCI bus and none expected; skipping GPU initialization (CPU-only runner)."
fi

# Step 1: Read JIT config from GCP instance metadata
log "Reading JIT config from instance metadata..."
METADATA_URL="http://metadata.google.internal/computeMetadata/v1/instance/attributes/jit-config"
MAX_RETRIES=10
JIT_CONFIG=""

for i in $(seq 1 "$MAX_RETRIES"); do
  JIT_CONFIG=$(curl -sf --max-time 10 --connect-timeout 5 -H "Metadata-Flavor: Google" "$METADATA_URL") && break
  log "  Attempt ${i}/${MAX_RETRIES}: Metadata not available yet, waiting..."
  sleep 5
done

if [ -z "$JIT_CONFIG" ]; then
  log "ERROR: Failed to read JIT config from metadata after $MAX_RETRIES attempts"
  shutdown -h now
  exit 1
fi

log "JIT config retrieved (${#JIT_CONFIG} chars)"

# Step 2: Log GPU and system info
log "=== System Information ==="
nvidia-smi 2>&1 | while read -r line; do log "  $line"; done || log "WARNING: nvidia-smi not available"
docker --version 2>&1 | while read -r line; do log "  $line"; done || log "WARNING: docker not available"

# Step 3: Run the GitHub Actions runner as the correct user
log "Starting runner as user '$RUNNER_USER' with JIT config..."
cd "$RUNNER_DIR"

# Run as the runner user, not root. The runner agent requires this.
EXIT_CODE=0
sudo -u "$RUNNER_USER" ./run.sh --jitconfig "$JIT_CONFIG" || EXIT_CODE=$?
log "Runner exited with code $EXIT_CODE"

# Step 4: Shut down the VM
log "=== Runner complete, shutting down VM ==="
shutdown -h now
