#!/usr/bin/env bash
# Install the pinned Lean 4 + Mathlib + REPL toolchain under /apps (Track L / P1).
#
# WHY this exists as a script and not as Python: the toolchain is an EXTERNAL BINARY
# (see docs/DEPENDENCY_POLICY.md, "Lean toolchain"). It is installed once, outside the
# Python environment, and `align/verify/lean` reaches it only through subprocess.
#
# Idempotent: every step is guarded by a marker/exists check, so re-running after a
# partial failure resumes instead of redoing a 5 GB download. Every network or build
# step runs under `timeout`, so a stalled mirror cannot hang the shell.
#
# Host rules (/apps/CLAUDE.md, binding): nothing is ever written to `/` or `/tmp`.
# elan, the Mathlib cache, git temp files and XDG caches are all redirected below.
# The preflight refuses on <45GB MemAvailable, load >= 2x cores, disk queue >= 64, or
# less than 40GB free on /apps. A Mathlib cache decompress is I/O heavy; it is not a
# GPU job so it does not take the GPU heavy-job lock.
#
# THE PIN (see docs/LEAN_VERIFIER.md section "Pin"):
#   Goedel-Prover-V2 follows DeepSeek-Prover-V1.5: its `mathlib4` git submodule is
#   xinhjBrant/mathlib4 @ 2f65ba7f1a9144b20c8e7358513548e317d26de1, a single commit on
#   top of leanprover-community/mathlib4 @ d1d1e4b72c6e009edd97cabe1d661b021fe8eeae
#   (2024-06-12) that ONLY edits lakefile.lean/lake-manifest.json to add DeepSeek's REPL
#   fork as a dependency. Its lean-toolchain is `leanprover/lean4:v4.9.0-rc1` (nominally
#   "Lean 4.9"; the binary reports 4.9.0-rc1). We pin the UPSTREAM parent: the Mathlib
#   .olean cache is keyed on file hashes that include lakefile.lean, so the DeepSeek
#   edit would orphan the cache, while the parent's .oleans are byte-for-byte the
#   Mathlib sources Goedel proofs were checked against. The REPL is upstream
#   leanprover-community/repl at the tag whose toolchain equals ours.
#
# Usage:  bash scripts/lean_setup.sh            # full install
#         LEAN_ROOT=... bash scripts/lean_setup.sh
#         bash scripts/lean_setup.sh --check    # only verify an existing install
set -euo pipefail

LEAN_ROOT="${LEAN_ROOT:-/apps/development/rlhf/lean}"
LEAN_TOOLCHAIN="${LEAN_TOOLCHAIN:-leanprover/lean4:v4.9.0-rc1}"
MATHLIB_REV="${MATHLIB_REV:-d1d1e4b72c6e009edd97cabe1d661b021fe8eeae}"
REPL_TAG="${REPL_TAG:-v4.9.0-rc1}"
REPL_URL="${REPL_URL:-https://github.com/leanprover-community/repl}"
MATHLIB_URL="https://github.com/leanprover-community/mathlib4"

case "$LEAN_ROOT" in
  /apps/*) ;;
  *) echo "refusing: LEAN_ROOT=$LEAN_ROOT is not under /apps (host rule: never write to / or /tmp)" >&2; exit 2 ;;
esac

export ELAN_HOME="$LEAN_ROOT/elan"
export MATHLIB_CACHE_DIR="$LEAN_ROOT/mathlib-cache"
export XDG_CACHE_HOME="$LEAN_ROOT/xdg-cache"
export XDG_CONFIG_HOME="$LEAN_ROOT/xdg-config"
export TMPDIR="$LEAN_ROOT/tmp"
export GIT_TERMINAL_PROMPT=0
export PATH="$ELAN_HOME/bin:$PATH"
WS="$LEAN_ROOT/workspace"
REPL_DIR="$LEAN_ROOT/repl"
LOGS="$LEAN_ROOT/logs"
mkdir -p "$LEAN_ROOT" "$MATHLIB_CACHE_DIR" "$XDG_CACHE_HOME" "$XDG_CONFIG_HOME" "$TMPDIR" "$LOGS"

log() { printf '[lean_setup %s] %s\n' "$(date +%H:%M:%S)" "$*"; }

preflight() {
  local avail_kb load cores q free_gb
  avail_kb=$(awk '/^MemAvailable:/{print $2}' /proc/meminfo)
  load=$(awk '{print $1}' /proc/loadavg)
  cores=$(nproc)
  q=$(awk '$3=="nvme2n1"{print $12}' /proc/diskstats)
  q=${q:-0}
  free_gb=$(df -BG --output=avail /apps | tail -1 | tr -dc 0-9)
  log "preflight: MemAvailable=$((avail_kb/1024/1024))GB load1=$load cores=$cores diskq=$q /apps_free=${free_gb}GB"
  [ "$avail_kb" -ge $((45*1024*1024)) ] || { log "REFUSED: MemAvailable < 45GB"; exit 83; }
  awk -v l="$load" -v c="$cores" 'BEGIN{exit !(l < 2*c)}' || { log "REFUSED: load >= 2x cores"; exit 83; }
  [ "$q" -lt 64 ] || { log "REFUSED: disk queue >= 64"; exit 83; }
  [ "$free_gb" -ge 40 ] || { log "REFUSED: < 40GB free on /apps"; exit 83; }
}

check_install() {
  [ -x "$ELAN_HOME/bin/lake" ] || { log "missing elan/lake"; return 1; }
  [ -x "$REPL_DIR/.lake/build/bin/repl" ] || { log "missing repl binary"; return 1; }
  [ -f "$WS/.lake/packages/mathlib/.lake/build/lib/Mathlib.olean" ] || { log "missing Mathlib.olean (cache not fetched)"; return 1; }
  log "install present"
}

if [ "${1:-}" = "--check" ]; then check_install; exit $?; fi

preflight

# ---- 1. elan -----------------------------------------------------------------------
if [ ! -x "$ELAN_HOME/bin/elan" ]; then
  log "installing elan into $ELAN_HOME"
  timeout 120 curl -sSfL https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -o "$LEAN_ROOT/elan-init.sh"
  # --no-modify-path: never touch ~/.profile / ~/.bashrc. --default-toolchain none: the
  # toolchain is installed explicitly below so the pin is visible in this file.
  timeout 600 sh "$LEAN_ROOT/elan-init.sh" -y --no-modify-path --default-toolchain none
fi
elan --version

# ---- 2. toolchain ------------------------------------------------------------------
if ! elan toolchain list | grep -qF "${LEAN_TOOLCHAIN##*:}"; then
  log "installing $LEAN_TOOLCHAIN"
  timeout 900 elan toolchain install "$LEAN_TOOLCHAIN"
fi

# ---- 3. workspace (Lake project requiring Mathlib at the pinned rev) ---------------
mkdir -p "$WS"
if [ ! -f "$WS/lakefile.lean" ]; then
  log "writing workspace lakefile"
  printf '%s\n' "$LEAN_TOOLCHAIN" > "$WS/lean-toolchain"
  cat > "$WS/lakefile.lean" <<EOF
import Lake
open Lake DSL

package «align_lean_ws»

require mathlib from git
  "$MATHLIB_URL" @ "$MATHLIB_REV"
EOF
fi
if [ ! -f "$WS/lake-manifest.json" ]; then
  log "lake update (clones mathlib + its pinned deps; several minutes)"
  (cd "$WS" && timeout 3000 lake update) 2>&1 | tee "$LOGS/lake_update.log" | tail -5
fi
# Guard against silent drift: the manifest must name exactly the pinned rev.
grep -q "$MATHLIB_REV" "$WS/lake-manifest.json" || { log "FATAL: manifest does not pin $MATHLIB_REV"; exit 3; }

# ---- 4. prebuilt Mathlib cache -----------------------------------------------------
if [ ! -f "$WS/.lake/packages/mathlib/.lake/build/lib/Mathlib.olean" ]; then
  log "lake exe cache get (downloads + unpacks the prebuilt oleans)"
  (cd "$WS" && timeout 3000 lake exe cache get) 2>&1 | tee "$LOGS/cache_get.log" | tail -8
fi
# Deps (batteries, Qq, aesop, proofwidgets) are built from source by lake when the cache
# does not cover them; this is a no-op when it does.
if ! (cd "$WS" && timeout 120 lake build --no-build Mathlib >/dev/null 2>&1); then
  log "lake build Mathlib (deps not fully cached; building the remainder)"
  (cd "$WS" && timeout 5400 lake build Mathlib) 2>&1 | tee "$LOGS/lake_build.log" | tail -5
fi

# ---- 5. REPL at the matching tag ---------------------------------------------------
if [ ! -d "$REPL_DIR/.git" ]; then
  log "cloning repl $REPL_TAG"
  timeout 300 git clone --branch "$REPL_TAG" --depth 1 "$REPL_URL" "$REPL_DIR"
fi
REPL_TC=$(cat "$REPL_DIR/lean-toolchain")
[ "$REPL_TC" = "$LEAN_TOOLCHAIN" ] || { log "FATAL: repl toolchain $REPL_TC != $LEAN_TOOLCHAIN"; exit 3; }
if [ ! -x "$REPL_DIR/.lake/build/bin/repl" ]; then
  log "building repl"
  (cd "$REPL_DIR" && timeout 1200 lake build) 2>&1 | tee "$LOGS/repl_build.log" | tail -5
fi

# ---- 6. smoke test: import Mathlib through the built REPL --------------------------
log "smoke test"
SMOKE=$(cd "$WS" && printf '%s\n\n%s\n\n' \
  '{"cmd": "import Mathlib\ntheorem t (a b : Nat) : a + b = b + a := by omega"}' \
  '{"cmd": "#print axioms t", "env": 0}' \
  | timeout 300 lake env "$REPL_DIR/.lake/build/bin/repl")
echo "$SMOKE" | head -20
echo "$SMOKE" | grep -q '"env"' || { log "FATAL: REPL smoke failed"; exit 4; }

cat > "$LEAN_ROOT/PINS.json" <<EOF
{"lean_toolchain": "$LEAN_TOOLCHAIN", "mathlib_rev": "$MATHLIB_REV",
 "repl_tag": "$REPL_TAG", "repl_rev": "$(git -C "$REPL_DIR" rev-parse HEAD)",
 "installed": "$(date -Iseconds)",
 "source_of_pin": "Goedel-Prover-V2 submodule xinhjBrant/mathlib4@2f65ba7f (parent leanprover-community/mathlib4@d1d1e4b7)"}
EOF
cat > "$LEAN_ROOT/env.sh" <<EOF
# source me. Redirects every cache the Lean toolchain writes to /apps.
export LEAN_ROOT="$LEAN_ROOT"
export ELAN_HOME="$ELAN_HOME"
export MATHLIB_CACHE_DIR="$MATHLIB_CACHE_DIR"
export XDG_CACHE_HOME="$XDG_CACHE_HOME"
export TMPDIR="$TMPDIR"
export PATH="\$ELAN_HOME/bin:\$PATH"
EOF
log "disk: $(du -sh "$LEAN_ROOT" 2>/dev/null | cut -f1) total under $LEAN_ROOT"
log "done"