#!/usr/bin/env bash
set -Eeuo pipefail

# Host launcher for Render Arena benchmark runs. Every prompt runs in a fresh
# container from a pinned benchmark image, with its reserved run directory
# mounted at /workspace and a temporary home holding only copied agent
# credentials. This file is also published as the standalone download.

SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)
IMAGE_PREFIX=${BENCH_IMAGE_PREFIX:-izolight/render-arena-bench}
IMAGE_TAG=${BENCH_IMAGE_TAG:-blender5.2.2}
RESULTS_ROOT=${BENCH_RESULTS_ROOT:-$SCRIPT_DIR/results}
GPU=false
[[ ${BENCH_GPU:-0} != 1 ]] || GPU=true
PROVIDER_ENV=(OPENAI_API_KEY ANTHROPIC_API_KEY OPENROUTER_API_KEY DEEPSEEK_API_KEY
  GEMINI_API_KEY GOOGLE_API_KEY ZAI_API_KEY)
PROXY_ENV=(HTTP_PROXY HTTPS_PROXY NO_PROXY http_proxy https_proxy no_proxy)
# Additional host variables passed to the next container (per action).
EXTRA_ENV=()
CATALOG=()

usage() {
  cat >&2 <<'EOF'
Usage:
  ./bench-run [run] TARGET --model PROVIDER/MODEL:EFFORT [--model ...] [--no-reasoning]
                    [--skip PROMPT_ID]... [--timeout DURATION] [--dry-run] [PROMPT_ID...]
  ./bench-run resume RUN_DIR [--timeout DURATION] [--prompt TEXT] [--dry-run]
  ./bench-run submit RUN_DIR [--dry-run|--yes]
  ./bench-run prompts [TARGET]

TARGET is pi, omp, omp-mcp, omp-mcp-skill, codex, or opencode. Each model and
prompt runs in its own fresh container; without prompt IDs, every catalog
prompt runs. --no-reasoning treats each --model value as an opaque model ID
and omits the reasoning level. --output-root DIR and --result-file FILE (one
model and prompt only) are used by the durable benchmark runner.

Environment: BENCH_RESULTS_ROOT (default: results next to this script),
BENCH_GPU=1 for Docker's --gpus all, RUN_TIMEOUT (default 45m, 0 disables),
BENCH_IMAGE_PREFIX (default izolight/render-arena-bench), BENCH_IMAGE_TAG,
RENDER_ARENA_URL and RENDER_ARENA_SUBMISSION_TOKEN for submissions.
Build and publish images from a checkout with containers/build-images.sh.
EOF
  exit 2
}

fail() {
  printf 'Error: %s\n' "$1" >&2
  exit "${2:-2}"
}

target_harness() {
  case $1 in omp-mcp|omp-mcp-skill) printf 'omp\n' ;; *) printf '%s\n' "$1" ;; esac
}

target_profile() {
  case $1 in omp-mcp) printf 'mcp\n' ;; omp-mcp-skill) printf 'mcp-skill\n' ;; *) printf 'normal\n' ;; esac
}

valid_target() {
  case $1 in pi|omp|omp-mcp|omp-mcp-skill|codex|opencode) return 0 ;; *) return 1 ;; esac
}

image_for() {
  printf '%s-%s:%s\n' "$IMAGE_PREFIX" "$1" "$IMAGE_TAG"
}

ensure_image() {
  local image=$1
  docker image inspect "$image" >/dev/null 2>&1 && return 0
  printf 'Pulling %s\n' "$image" >&2
  docker pull "$image" >&2 || fail "benchmark image is unavailable: $image (build it with containers/build-images.sh)" 1
}

# Rootless Docker maps container UID 0 to the daemon account on the host.
# Reusing the host UID instead maps it into the subordinate ID range, which
# cannot access bind mounts owned by the daemon account.
CONTAINER_USER=
resolve_container_user() {
  [[ -z $CONTAINER_USER ]] || return 0
  local security_options
  CONTAINER_USER="$(id -u):$(id -g)"
  if security_options=$(docker info --format '{{json .SecurityOptions}}' 2>/dev/null) \
    && [[ $security_options == *'"name=rootless"'* ]]; then
    CONTAINER_USER=0:0
  fi
}

copy_file() {
  local source=$1 destination=$2
  [[ -f $source ]] || return 0
  mkdir -p -- "${destination%/*}"
  cp -- "$source" "$destination"
  chmod 600 "$destination"
}

copy_sqlite() {
  local source=$1 destination=$2
  [[ -f $source ]] || return 0
  mkdir -p -- "${destination%/*}"
  python3 - "$source" "$destination" <<'PY'
import sqlite3, sys
source = sqlite3.connect(f"file:{sys.argv[1]}?mode=ro", uri=True)
destination = sqlite3.connect(sys.argv[2])
source.backup(destination)
destination.close()
source.close()
PY
  chmod 600 "$destination"
}

# Copy only the credentials and settings each harness needs into the temporary
# container home. Nothing from it is copied into the results tree.
stage_agent_home() {
  local harness=$1 home=$2
  local data=${XDG_DATA_HOME:-$HOME/.local/share} config=${XDG_CONFIG_HOME:-$HOME/.config}
  case $harness in
    pi)
      copy_file "$HOME/.pi/agent/auth.json" "$home/.pi/agent/auth.json"
      copy_file "$HOME/.pi/agent/models.json" "$home/.pi/agent/models.json"
      copy_file "$HOME/.pi/agent/settings.json" "$home/.pi/agent/settings.json"
      ;;
    omp)
      copy_sqlite "$HOME/.omp/agent/agent.db" "$home/.omp/agent/agent.db"
      copy_sqlite "$HOME/.omp/agent/models.db" "$home/.omp/agent/models.db"
      copy_file "$HOME/.omp/agent/auth.json" "$home/.omp/agent/auth.json"
      copy_file "$HOME/.omp/agent/models.json" "$home/.omp/agent/models.json"
      copy_file "$HOME/.omp/agent/settings.json" "$home/.omp/agent/settings.json"
      copy_file "$HOME/.omp/config.yml" "$home/.omp/config.yml"
      copy_file "$HOME/.omp/install-id" "$home/.omp/install-id"
      ;;
    codex)
      copy_file "$HOME/.codex/auth.json" "$home/.codex/auth.json"
      copy_file "$HOME/.codex/config.toml" "$home/.codex/config.toml"
      ;;
    opencode)
      copy_sqlite "$data/opencode/opencode.db" "$home/.local/share/opencode/opencode.db"
      copy_file "$data/opencode/auth.json" "$home/.local/share/opencode/auth.json"
      copy_file "$config/opencode/account.json" "$home/.config/opencode/account.json"
      copy_file "$config/opencode/opencode.json" "$home/.config/opencode/opencode.json"
      copy_file "$config/opencode/opencode.jsonc" "$home/.config/opencode/opencode.jsonc"
      ;;
  esac
}

# The temporary directory of the active container. The traps remove its
# container and credentials when a run finishes or the launcher is stopped.
ACTIVE_TEMP=
remove_active_container() {
  [[ -n $ACTIVE_TEMP ]] || return 0
  if [[ -s $ACTIVE_TEMP/container.cid ]]; then
    docker rm -f "$(<"$ACTIVE_TEMP/container.cid")" >/dev/null 2>&1 || true
  fi
  rm -rf -- "$ACTIVE_TEMP"
  ACTIVE_TEMP=
}
trap remove_active_container EXIT
trap 'exit 130' INT
trap 'exit 143' TERM HUP

# Run IMAGE with RUN_DIR at /workspace. HARNESS selects the staged
# credentials; an empty value starts with an empty home.
container_run() {
  local image=$1 run_dir=$2 harness=$3
  shift 3
  local status=0 name
  resolve_container_user
  ACTIVE_TEMP=$(mktemp -d)
  chmod 700 "$ACTIVE_TEMP"
  mkdir -m 700 "$ACTIVE_TEMP/home"
  [[ -z $harness ]] || stage_agent_home "$harness" "$ACTIVE_TEMP/home"
  local -a args=(run --rm -i --init --cidfile "$ACTIVE_TEMP/container.cid" --user "$CONTAINER_USER"
    -v "$run_dir:/workspace" -v "$ACTIVE_TEMP/home:/home/agent")
  [[ -t 0 && -t 1 ]] && args+=(-t)
  [[ $GPU == false ]] || args+=(--gpus all -e BENCH_GPU=1)
  for name in "${PROXY_ENV[@]}" "${EXTRA_ENV[@]}"; do
    [[ ${!name+x} ]] && args+=(-e "$name")
  done
  docker "${args[@]}" "$image" "$@" || status=$?
  remove_active_container
  return "$status"
}

# Prompt IDs come from the image, which is the source of truth for the catalog.
load_prompts() {
  local image=$1
  mapfile -t CATALOG < <(docker run --rm "$image" prompts)
  ((${#CATALOG[@]} > 0)) || fail "could not read the prompt catalog from $image" 1
}

known_prompt() {
  local prompt
  for prompt in "${CATALOG[@]}"; do
    [[ $prompt == "$1" ]] && return 0
  done
  return 1
}

check_component() {
  local label=$1 value=$2
  [[ -n $value && $value != . && $value != .. && $value != */* && $value != *[[:cntrl:]]* ]] \
    || fail "invalid $label: $value"
}

json_string() {
  local value=$1
  value=${value//\\/\\\\}
  value=${value//\"/\\\"}
  printf '"%s"' "$value"
}

# mkdir without -p is atomic, so concurrent launchers never share a run.
reserve_run_dir() {
  local base=$1 index=0 path name
  mkdir -p -- "$base" || return 1
  shopt -s nullglob
  for path in "$base"/*/; do
    name=${path%/}
    name=${name##*/}
    [[ $name =~ ^[0-9]+$ ]] || continue
    ((10#$name < index)) || index=$((10#$name + 1))
  done
  shopt -u nullglob
  while ! mkdir -- "$base/$index" 2>/dev/null; do
    # A collision is expected when another launcher reserves the same index.
    # Any other failure (especially permissions) must stop the run.
    if [[ ! -e $base/$index ]]; then
      printf 'Error: could not reserve run directory under %s; check its ownership and ACLs.\n' "$base" >&2
      return 1
    fi
    index=$((index + 1))
  done
  printf '%s\n' "$base/$index"
}

# Run one prompt in a fresh container below OUTPUT_ROOT/MODEL/EFFORT/PROMPT/N.
# Sets RUN_DIR to the reserved host directory.
RUN_DIR=
run_one() {
  local target=$1 output_root=$2 model_spec=$3 effort=$4 prompt=$5 timeout=$6 dry_run=$7
  local harness model_dir
  harness=$(target_harness "$target")
  [[ $model_spec == */* && -n ${model_spec%%/*} && -n ${model_spec#*/} ]] \
    || fail "model must be PROVIDER/MODEL: $model_spec"
  model_dir=${model_spec#*/}
  model_dir=${model_dir##*/}
  check_component model "$model_dir"
  check_component effort "$effort"
  check_component prompt "$prompt"
  local -a args=(run --model "$model_spec" --prompt "$prompt")
  if [[ $effort == none ]]; then args+=(--no-reasoning); else args+=(--effort "$effort"); fi
  [[ -z $timeout ]] || args+=(--timeout "$timeout")
  if [[ $dry_run == true ]]; then
    local scratch
    scratch=$(mktemp -d)
    RUN_DIR=
    container_run "$(image_for "$target")" "$scratch" "" "${args[@]}" --dry-run || { rm -rf -- "$scratch"; return 1; }
    rm -rf -- "$scratch"
    return 0
  fi
  RUN_DIR=$(reserve_run_dir "$output_root/$model_dir/$effort/$prompt") || return 1
  mkdir -p -- "$RUN_DIR/.benchmark"
  printf '{\n  "harness": %s,\n  "profile": %s,\n  "model": %s,\n  "effort": %s,\n  "prompt": %s\n}\n' \
    "$(json_string "$harness")" "$(json_string "${output_root##*/}")" "$(json_string "$model_spec")" \
    "$(json_string "$effort")" "$(json_string "$prompt")" >"$RUN_DIR/.benchmark/run.json"
  printf 'Run directory: %s\n' "$RUN_DIR"
  EXTRA_ENV=("${PROVIDER_ENV[@]}")
  container_run "$(image_for "$target")" "$RUN_DIR" "$harness" "${args[@]}"
}

results_root() {
  mkdir -p -- "$RESULTS_ROOT"
  (cd -- "$RESULTS_ROOT" && pwd -P)
}

# Split RUN_DIR below the results root into harness/profile/model/effort/prompt/run.
RUN_PARTS=()
existing_run() {
  local root run_dir relative
  root=$(results_root)
  [[ -d $1 ]] || fail "run directory does not exist: $1" 1
  run_dir=$(cd -- "$1" && pwd -P)
  case $run_dir in "$root"/*) ;; *) fail "run must be inside $root" ;; esac
  relative=${run_dir#"$root"/}
  IFS=/ read -r -a RUN_PARTS <<<"$relative"
  ((${#RUN_PARTS[@]} == 6)) || fail "invalid run directory layout: $run_dir"
  RUN_DIR=$run_dir
}

run_target() {
  local harness=$1 profile=$2
  case $harness:$profile in
    pi:normal|omp:normal|codex:normal|opencode:normal) printf '%s\n' "$harness" ;;
    omp:mcp) printf 'omp-mcp\n' ;;
    omp:mcp-skill) printf 'omp-mcp-skill\n' ;;
    *) return 1 ;;
  esac
}

cmd_run() {
  (($# >= 1)) || usage
  local target=$1
  shift
  valid_target "$target" || usage
  local -a models=() skip=() requested=()
  local no_reasoning=false timeout=${RUN_TIMEOUT:-} dry_run=false output_root= result_file=
  while (($#)); do
    case $1 in
      --model) models+=("${2:?--model needs PROVIDER/MODEL:EFFORT}"); shift 2 ;;
      --model=*) models+=("${1#*=}"); shift ;;
      --no-reasoning|--no-thinking|--no-reasoning-level) no_reasoning=true; shift ;;
      --skip) skip+=("${2:?--skip needs a prompt ID}"); shift 2 ;;
      --skip=*) skip+=("${1#*=}"); shift ;;
      --timeout) timeout=${2:?--timeout needs a duration}; shift 2 ;;
      --timeout=*) timeout=${1#*=}; shift ;;
      --output-root) output_root=${2:?--output-root needs a directory}; shift 2 ;;
      --result-file) result_file=${2:?--result-file needs a path}; shift 2 ;;
      --dry-run) dry_run=true; shift ;;
      --) shift; requested+=("$@"); break ;;
      -*) fail "unknown run option: $1" ;;
      *) requested+=("$1"); shift ;;
    esac
  done
  ((${#models[@]} > 0)) || fail 'at least one --model PROVIDER/MODEL:EFFORT is required.'
  local model
  local -a specs=() efforts=()
  for model in "${models[@]}"; do
    if [[ $no_reasoning == true ]]; then
      # The value is an opaque model identifier; a colon is not an effort.
      specs+=("$model"); efforts+=(none)
    else
      [[ $model == *:* ]] || fail "--model must be PROVIDER/MODEL:EFFORT (or use --no-reasoning): $model"
      specs+=("${model%:*}"); efforts+=("${model##*:}")
    fi
  done

  local image id prompt selected
  local -a prompts=()
  image=$(image_for "$target")
  ensure_image "$image"
  load_prompts "$image"
  for id in "${requested[@]}" "${skip[@]}"; do
    known_prompt "$id" || fail "unknown prompt ID: $id"
  done
  # Without prompt IDs, run the whole catalog in catalog order.
  if ((${#requested[@]} == 0)); then requested=("${CATALOG[@]}"); fi
  for prompt in "${requested[@]}"; do
    selected=true
    for id in "${skip[@]}"; do [[ $id != "$prompt" ]] || selected=false; done
    [[ $selected == false ]] || prompts+=("$prompt")
  done
  ((${#prompts[@]} > 0)) || fail 'no prompts selected.'

  if [[ -n $result_file ]]; then
    ((${#specs[@]} == 1 && ${#prompts[@]} == 1)) && [[ $dry_run == false ]] \
      || fail '--result-file requires exactly one model and one prompt, and no --dry-run.'
    mkdir -p -- "$(dirname -- "$result_file")"
  fi
  [[ -n $output_root ]] || output_root="$(results_root)/$(target_harness "$target")/$(target_profile "$target")"
  mkdir -p -- "$output_root"
  output_root=$(cd -- "$output_root" && pwd -P)

  local spec failed=0 total=0
  for spec in "${!specs[@]}"; do
    for prompt in "${prompts[@]}"; do
      total=$((total + 1))
      run_one "$target" "$output_root" "${specs[spec]}" "${efforts[spec]}" "$prompt" "$timeout" "$dry_run" \
        || failed=$((failed + 1))
    done
  done
  ((failed == 0)) || fail "$failed of $total run(s) failed or timed out." 1
  # Publish the run directory only after success; the rename keeps readers
  # (the durable runner) from observing a partial result file.
  if [[ -n $result_file ]]; then
    printf '%s\n' "$RUN_DIR" >"$result_file.tmp.$$"
    mv -f -- "$result_file.tmp.$$" "$result_file"
    printf 'Result file: %s\n' "$result_file"
  fi
}

cmd_resume() {
  (($# >= 1)) || usage
  existing_run "$1"
  shift
  local target image
  target=$(run_target "${RUN_PARTS[0]}" "${RUN_PARTS[1]}") \
    || fail "no benchmark image resumes ${RUN_PARTS[0]}/${RUN_PARTS[1]} runs"
  image=$(image_for "$target")
  ensure_image "$image"
  EXTRA_ENV=("${PROVIDER_ENV[@]}" RUN_TIMEOUT)
  container_run "$image" "$RUN_DIR" "${RUN_PARTS[0]}" resume "$@"
}

cmd_submit() {
  (($# >= 1)) || usage
  existing_run "$1"
  shift
  local arg target image
  for arg in "$@"; do
    case $arg in --dry-run|--yes) ;; *) usage ;; esac
  done
  target=$(run_target "${RUN_PARTS[0]}" "${RUN_PARTS[1]}") \
    || fail "no benchmark image submits ${RUN_PARTS[0]}/${RUN_PARTS[1]} runs; use static/submit-run.sh --dir"
  image=$(image_for "$target")
  ensure_image "$image"
  # Runs reserved by this launcher record their metadata in
  # .benchmark/run.json, which the image reads. Older runs use their path.
  local -a metadata=()
  if [[ ! -f $RUN_DIR/.benchmark/run.json ]]; then
    metadata=(--harness "${RUN_PARTS[0]}" --effort "${RUN_PARTS[3]}" --prompt "${RUN_PARTS[4]}")
  fi
  EXTRA_ENV=(RENDER_ARENA_URL RENDER_ARENA_SUBMISSION_TOKEN)
  container_run "$image" "$RUN_DIR" "" submit "${metadata[@]}" "$@"
}

case ${1:-} in
  run) shift; cmd_run "$@" ;;
  resume) shift; cmd_resume "$@" ;;
  submit) shift; cmd_submit "$@" ;;
  prompts)
    shift
    (($# <= 1)) || usage
    target=${1:-pi}
    valid_target "$target" || usage
    ensure_image "$(image_for "$target")"
    load_prompts "$(image_for "$target")"
    printf '%s\n' "${CATALOG[@]}"
    ;;
  *) valid_target "${1:-}" || usage; cmd_run "$@" ;;
esac