#!/usr/bin/env bash
# Keep a ready Lean + Mathlib project on the runner host, keyed by the project's
# toolchain and lakefile, so only the first job pays for the multi-GB cache fetch.
#
#   prepare-mathlib.sh <lean-project-dir>     prints the cached project path
#
# The runner executes jobs on the host (no container), so /opt/lean-cache persists
# between jobs. A flock serialises concurrent jobs; .ready is written last, so an
# interrupted preparation is redone rather than trusted.
set -euo pipefail

SRC="${1:?usage: prepare-mathlib.sh <lean-project-dir>}"
ROOT="${LEAN_CACHE_ROOT:-/opt/lean-cache}"
KEY="$(cat "$SRC/lean-toolchain" "$SRC/lakefile.toml" | sha256sum | cut -c1-12)"
DIR="$ROOT/leanbench-$KEY"

mkdir -p "$ROOT"
exec 9>"$ROOT/.lock"
flock 9

if [ -f "$DIR/.ready" ]; then
  echo "prepare-mathlib: reusing $DIR" >&2
else
  echo "prepare-mathlib: building $DIR (first run on this toolchain/lakefile)" >&2
  rm -rf "$DIR"
  cp -r "$SRC" "$DIR"
  cd "$DIR"
  elan toolchain install "$(cat lean-toolchain)" >&2
  # Keep the cache fetch out of the post-update hook so it goes through the retries.
  MATHLIB_NO_CACHE_ON_UPDATE=1 timeout 1800 lake update >&2
  # The cache is ~9,000 files; a few transient transfer failures are normal, and a
  # re-run fetches only what is still missing (downloads persist in ~/.cache/mathlib).
  for attempt in 1 2 3 4 5 6; do
    if timeout 5400 lake exe cache get >&2; then break; fi
    [ "$attempt" -eq 6 ] && { echo "prepare-mathlib: cache get failed 6 times" >&2; exit 1; }
    echo "prepare-mathlib: cache get attempt $attempt failed; retrying" >&2
    sleep $((attempt * 15))
  done
  start=$(date +%s)
  timeout 1200 lake env lean LeanBench.lean >&2 || {
    status=$?
    [ "$status" -eq 124 ] && echo "prepare-mathlib: 'import Mathlib' timed out; check RAM (free -h)" >&2
    exit "$status"
  }
  echo "prepare-mathlib: import Mathlib elaborated in $(( $(date +%s) - start ))s" >&2
  touch .ready
fi
echo "$DIR"