#!/usr/bin/env bash
# Replay harness: Lean proof check of ONE helper declaration from the public
# Palomar all-dimensions source.
#
#   module      SpectreFinite.RadiusFiveRoot8Step1Chunk8_2Slice3
#   declaration SpectreFinite.HexSearch.RadiusFive.root8Node0Step1Chunk8_2_slice3_checked
#   commit      048bb245b60abce79f1a2d0c82336ff8508ec752
#
# What it does:
#   1. verifies the checkout is byte-identical to the pinned commit for src/,
#      lean-toolchain, lakefile.lean and lake-manifest.json (nothing is edited);
#   2. obtains the project-pinned toolchain leanprover/lean4:v4.35.0-rc2
#      (official release tarball, sha256-verified) unless LEAN_HOME is given;
#   3. computes the transitive import closure of the module (it imports only
#      project modules and Lean core `Init`, so Mathlib and the other Lake
#      packages are pinned but never loaded) and compiles every module of
#      that closure from source with the library's moreLeanArgs
#      (-j1 -DElab.async=false), in dependency order;
#   4. elaborates an audit file that (a) re-states the theorem's exact type and
#      checks the declaration against it, (b) runs `#print axioms`, and
#      (c) collects the transitive axioms programmatically and fails unless they
#      are a subset of {propext, Classical.choice, Quot.sound};
#   5. writes evidence.txt and receipt.json next to this script (or OUT_DIR).
#
# All Lean invocations (steps 3 and 4) share ONE wall-clock budget of
# CHECK_TIMEOUT seconds (default and maximum 1200). If it expires the result
# is INCONCLUSIVE, never PASS.
#
# Environment knobs (all optional):
#   REPO          repository root            (default: git toplevel of this script)
#   OUT_DIR       where evidence/receipt go  (default: directory of this script)
#   WORK_DIR      scratch dir for toolchain/oleans/logs (default: mktemp -d)
#   LEAN_HOME     existing toolchain dir containing bin/lean (skips download)
#   JOBS          parallel module compilations (default: nproc)
#   CHECK_TIMEOUT seconds, <= 1200           (default: 1200)
#
# Exit status: 0 PASS, 1 FAIL, 2 INCONCLUSIVE (timeout), 3 setup/pin error.
set -uo pipefail

EXPECTED_COMMIT=048bb245b60abce79f1a2d0c82336ff8508ec752
MODULE=SpectreFinite.RadiusFiveRoot8Step1Chunk8_2Slice3
DECL=SpectreFinite.HexSearch.RadiusFive.root8Node0Step1Chunk8_2_slice3_checked
TOOLCHAIN=leanprover/lean4:v4.35.0-rc2
LEAN_VERSION=4.35.0-rc2
LEAN_GITHASH=11acb17ec6b07a8f9e9173e6845197929540936b
LEAN_URL=https://github.com/leanprover/lean4/releases/download/v4.35.0-rc2/lean-4.35.0-rc2-linux.tar.zst
LEAN_SHA256=b67d0e0ba959fc74e186f1b44aebca32d6ecb6a28f8e9a029fc4e9f73398a049
MATHLIB_REV=065356127b1dc0016f66b7283ce0ce2c4055aa55
LEAN_ARGS=(-j1 -DElab.async=false)   # lakefile.lean: lean_lib PalomarAllDimensions moreLeanArgs
ALLOWED_AXIOMS="Classical.choice propext Quot.sound"

HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
REPO=${REPO:-$(git -C "$HERE" rev-parse --show-toplevel 2>/dev/null)}
OUT_DIR=${OUT_DIR:-$HERE}
WORK_DIR=${WORK_DIR:-$(mktemp -d "${TMPDIR:-/tmp}/lean-proof-unit.XXXXXX")}
JOBS=${JOBS:-$(nproc 2>/dev/null || echo 4)}
CHECK_TIMEOUT=${CHECK_TIMEOUT:-1200}
if ! [[ "$CHECK_TIMEOUT" =~ ^[0-9]+$ ]] || (( CHECK_TIMEOUT < 1 || CHECK_TIMEOUT > 1200 )); then
  echo "CHECK_TIMEOUT must be an integer in 1..1200" >&2; exit 3
fi

EVIDENCE="$OUT_DIR/evidence.txt"
RECEIPT="$OUT_DIR/receipt.json"
mkdir -p "$OUT_DIR" "$WORK_DIR"
: > "$EVIDENCE"
ev()  { printf '%s\n' "$*" >> "$EVIDENCE"; }
log() { printf '[replay] %s\n' "$*" >&2; ev "$*"; }
die() { log "SETUP ERROR: $*"; write_receipt "ERROR" 3 "$*"; exit 3; }
now() { date -u +%Y-%m-%dT%H:%M:%SZ; }

RESULT=""; CHECK_EXIT=""; AXIOMS_JSON="null"; MODULE_COUNT=0; ELAPSED=""
SRC_SHA=""; TC_SHA=""; LF_SHA=""; LM_SHA=""; HEAD_COMMIT=""; LEAN_VERSION_LINE=""
write_receipt() {  # status exit reason
  python3 -I - "$RECEIPT" "$1" "$2" "$3" <<PY
import json, sys
path, status, code, reason = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]
r = {
  "schema": "lean-proof-unit/receipt/v1",
  "result": status,
  "result_reason": reason,
  "harness_exit_status": code,
  "lean_check_exit_status": None if "$CHECK_EXIT" == "" else int("$CHECK_EXIT"),
  "checked_at_utc": "$(now)",
  "source": {
    "repository_head": "$HEAD_COMMIT" or None,
    "pinned_commit": "$EXPECTED_COMMIT",
    "tree_matches_pinned_commit_for": ["src/", "lean-toolchain", "lakefile.lean", "lake-manifest.json"],
    "module": "$MODULE",
    "module_file": "src/${MODULE//.//}.lean",
    "module_file_sha256": "$SRC_SHA" or None,
    "declaration": "$DECL",
  },
  "pins": {
    "lean_toolchain_file": "$TOOLCHAIN",
    "lean_toolchain_file_sha256": "$TC_SHA" or None,
    "lean_release_url": "$LEAN_URL",
    "lean_release_sha256": "$LEAN_SHA256",
    "lean_version_reported": "$LEAN_VERSION_LINE" or None,
    "lakefile_lean_sha256": "$LF_SHA" or None,
    "lake_manifest_json_sha256": "$LM_SHA" or None,
    "mathlib_rev": "$MATHLIB_REV",
    "mathlib_loaded_by_this_check": False,
    "lean_args": ["-j1", "-DElab.async=false"],
  },
  "check": {
    "method": "lean -o (direct compile of transitive import closure, then audit file)",
    "modules_compiled": $MODULE_COUNT,
    "timeout_seconds": $CHECK_TIMEOUT,
    "elapsed_seconds": None if "$ELAPSED" == "" else int("$ELAPSED"),
    "allowed_axioms": ["propext", "Classical.choice", "Quot.sound"],
    "axioms_reported": json.loads('$AXIOMS_JSON'),
  },
  "not_performed": [
    "final all-dimensional theorem comparison (comparator Challenge/Solution)",
    "independent-kernel replay (e.g. lean4checker / external kernel)",
    "Lake build of the full PalomarAllDimensions library",
    "official Palomar preflight",
  ],
  "scope_note": "Checks one helper declaration only. It does not establish the all-dimensional theorem or Palomar verification.",
  "evidence_file": "evidence.txt",
}
json.dump(r, open(path, "w"), indent=2); open(path, "a").write("\n")
PY
}

ev "# Lean proof-check evidence (generated by replay.sh; do not edit by hand)"
ev "started_utc: $(now)"
ev "module: $MODULE"
ev "declaration: $DECL"
ev "work_dir: (scratch, not delivered)"
ev ""

# ---------------------------------------------------------------- 1. pins
[[ -n "$REPO" && -d "$REPO/src" ]] || die "cannot locate repository root (set REPO)"
cd "$REPO" || die "cd $REPO"
HEAD_COMMIT=$(git rev-parse HEAD 2>/dev/null)
git cat-file -e "$EXPECTED_COMMIT^{commit}" 2>/dev/null || die "pinned commit $EXPECTED_COMMIT not present in repository"
git diff --quiet "$EXPECTED_COMMIT" -- src lean-toolchain lakefile.lean lake-manifest.json \
  || die "working tree differs from pinned commit in src/ or build pins"
[[ -z "$(git ls-files --others --exclude-standard -- src)" ]] || die "untracked files under src/"
[[ "$(cat lean-toolchain)" == "$TOOLCHAIN" ]] || die "lean-toolchain is not $TOOLCHAIN"
python3 -I -c 'import json,sys; m=json.load(open("lake-manifest.json")); r=[p["rev"] for p in m["packages"] if p["name"]=="mathlib"]; sys.exit(0 if r==[sys.argv[1]] else 1)' "$MATHLIB_REV" \
  || die "lake-manifest.json mathlib rev is not $MATHLIB_REV"
MODFILE="src/${MODULE//.//}.lean"
[[ -f "$MODFILE" ]] || die "missing $MODFILE"
SRC_SHA=$(sha256sum "$MODFILE" | cut -d' ' -f1)
TC_SHA=$(sha256sum lean-toolchain | cut -d' ' -f1)
LF_SHA=$(sha256sum lakefile.lean | cut -d' ' -f1)
LM_SHA=$(sha256sum lake-manifest.json | cut -d' ' -f1)
log "== 1. source and pins"
log "repository HEAD: $HEAD_COMMIT"
log "pinned commit:   $EXPECTED_COMMIT (src/, lean-toolchain, lakefile.lean, lake-manifest.json identical: yes)"
log "lean-toolchain:  $TOOLCHAIN  sha256 $TC_SHA"
log "lakefile.lean    sha256 $LF_SHA"
log "lake-manifest    sha256 $LM_SHA  (mathlib $MATHLIB_REV)"
log "$MODFILE sha256 $SRC_SHA"
ev ""

# ---------------------------------------------------------------- 2. toolchain
log "== 2. toolchain"
if [[ -z "${LEAN_HOME:-}" ]]; then
  log "\$ curl -sSLf -o lean.tar.zst $LEAN_URL"
  curl -sSLf -o "$WORK_DIR/lean.tar.zst" "$LEAN_URL" || die "toolchain download failed"
  got=$(sha256sum "$WORK_DIR/lean.tar.zst" | cut -d' ' -f1)
  log "sha256(lean.tar.zst) = $got (expected $LEAN_SHA256)"
  [[ "$got" == "$LEAN_SHA256" ]] || die "toolchain sha256 mismatch"
  tar --zstd -xf "$WORK_DIR/lean.tar.zst" -C "$WORK_DIR" || die "toolchain extract failed"
  LEAN_HOME="$WORK_DIR/lean-$LEAN_VERSION-linux"
else
  log "using pre-installed LEAN_HOME (must be the official $LEAN_VERSION release; version and githash verified below)"
fi
LEAN="$LEAN_HOME/bin/lean"
[[ -x "$LEAN" ]] || die "no lean binary at $LEAN"
LEAN_VERSION_LINE=$("$LEAN" --version)
log "\$ lean --version"
log "$LEAN_VERSION_LINE"
[[ "$LEAN_VERSION_LINE" == *"version $LEAN_VERSION,"*"commit $LEAN_GITHASH"* ]] || die "lean is not $LEAN_VERSION ($LEAN_GITHASH)"
ev ""

# ---------------------------------------------------------------- 3+4. check
BUILD="$WORK_DIR/build"; LOGS="$WORK_DIR/logs"; AUDIT_DIR="$WORK_DIR/audit"
mkdir -p "$BUILD" "$LOGS" "$AUDIT_DIR"
AUDIT="$AUDIT_DIR/LeanProofUnitAudit.lean"
cat > "$AUDIT" <<'LEAN'
import SpectreFinite.RadiusFiveRoot8Step1Chunk8_2Slice3
import Lean

-- same namespace and `open`s as the published source file
namespace SpectreFinite.HexSearch.RadiusFive
open SpectreFinite.FiniteTable
open HigherFastChecking

-- (a) the declaration has exactly the statement written in the published source
example :
    checkNarrowing ⟨table56Scope, RadiusFourPruning.stateRows8_2_3⟩
      root8Node0Domains0 root8Node0Domains1 = true :=
  SpectreFinite.HexSearch.RadiusFive.root8Node0Step1Chunk8_2_slice3_checked

#check @SpectreFinite.HexSearch.RadiusFive.root8Node0Step1Chunk8_2_slice3_checked

-- (b) human-readable transitive axiom report
#print axioms SpectreFinite.HexSearch.RadiusFive.root8Node0Step1Chunk8_2_slice3_checked

-- (c) machine audit: must be a theorem; axioms must be within the allowed set
open Lean Elab Command in
run_cmd do
  let n := ``SpectreFinite.HexSearch.RadiusFive.root8Node0Step1Chunk8_2_slice3_checked
  let env ← getEnv
  let some ci := env.find? n | throwError "AUDIT FAIL: {n} not found"
  unless ci matches .thmInfo _ do throwError "AUDIT FAIL: {n} is not a theorem"
  let axs ← liftCoreM (Lean.collectAxioms n)
  let allowed : List Name := [``propext, ``Classical.choice, ``Quot.sound]
  let bad := axs.toList.filter (fun a => !allowed.contains a)
  logInfo m!"AUDIT_AXIOMS: {axs.toList}"
  unless bad.isEmpty do throwError "AUDIT FAIL: disallowed axioms {bad}"
  logInfo m!"AUDIT OK: {n} is a theorem; axioms within [propext, Classical.choice, Quot.sound]"

end SpectreFinite.HexSearch.RadiusFive
LEAN

SCHED="$WORK_DIR/sched.py"
cat > "$SCHED" <<'PY'
# Compile the transitive import closure of one module with `lean -o`, in
# dependency order, then elaborate the audit file. Prints one line per module.
import os, re, subprocess, sys, time, concurrent.futures as cf
lean, src, build, logs, target, audit, jobs = sys.argv[1:8]
extra = sys.argv[8:]
core = {"Init", "Std", "Lean", "Lake"}
imp = re.compile(r"^\s*(?:public\s+)?(?:meta\s+)?import\s+(?:all\s+)?(.+)$")
def path(m): return os.path.join(src, *m.split(".")) + ".lean"
deps, stack = {}, [target]
while stack:
    m = stack.pop()
    if m in deps: continue
    p = path(m)
    if not os.path.exists(p):
        print(f"MISSING-MODULE {m}", flush=True); sys.exit(3)
    ds = []
    for line in open(p, encoding="utf-8"):
        r = imp.match(line)
        if r:
            for d in r.group(1).split():
                if d.split(".")[0] in core: continue
                if not os.path.exists(path(d)):
                    print(f"NON-LOCAL-IMPORT {m} -> {d} (would need Lake packages)", flush=True); sys.exit(3)
                ds.append(d)
    deps[m] = ds; stack += ds
print(f"CLOSURE {len(deps)} project modules (all other imports are Lean core)", flush=True)
env = dict(os.environ, LEAN_PATH=build)
def compile_(m):
    base = os.path.join(build, *m.split("."))
    os.makedirs(os.path.dirname(base), exist_ok=True)
    cmd = [lean, "-R", src, "-o", base + ".olean", "-i", base + ".ilean", *extra, path(m)]
    t = time.time()
    r = subprocess.run(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    open(os.path.join(logs, m + ".log"), "wb").write(r.stdout)
    return m, r.returncode, time.time() - t
done, failed, running = set(), [], {}
with cf.ThreadPoolExecutor(int(jobs)) as ex:
    while len(done) < len(deps) and not failed:
        for m, ds in deps.items():
            if m not in done and m not in running and all(d in done for d in ds):
                running[m] = ex.submit(compile_, m)
        fin, _ = cf.wait(running.values(), return_when=cf.FIRST_COMPLETED)
        for f in fin:
            m, rc, dt = f.result(); del running[m]
            print(f"MODULE {m} exit={rc} seconds={dt:.1f}", flush=True)
            (done.add(m) if rc == 0 else failed.append(m))
if failed:
    print(f"COMPILE-FAILED {failed}", flush=True); sys.exit(1)
cmd = [lean, *extra, audit]
print("AUDIT-COMMAND LEAN_PATH=<build> " + " ".join(["lean", *extra, os.path.basename(audit)]), flush=True)
r = subprocess.run(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
open(os.path.join(logs, "audit.log"), "wb").write(r.stdout)
print(f"AUDIT exit={r.returncode}", flush=True)
sys.exit(0 if r.returncode == 0 else 1)
PY

log "== 3. compile transitive closure + 4. audit (single budget: timeout ${CHECK_TIMEOUT}s)"
CMD=(timeout --kill-after=10 "$CHECK_TIMEOUT" python3 -I "$SCHED" "$LEAN" "$REPO/src" "$BUILD" "$LOGS" "$MODULE" "$AUDIT" "$JOBS" "${LEAN_ARGS[@]}")
log "\$ timeout --kill-after=10 $CHECK_TIMEOUT python3 -I sched.py <lean> src <build> <logs> $MODULE LeanProofUnitAudit.lean $JOBS ${LEAN_ARGS[*]}"
log "  per module: LEAN_PATH=<build> lean -R src -o <build>/M.olean -i <build>/M.ilean ${LEAN_ARGS[*]} src/M.lean"
T0=$(date +%s)
"${CMD[@]}" > "$WORK_DIR/sched.out" 2>&1
CHECK_EXIT=$?
ELAPSED=$(( $(date +%s) - T0 ))
MODULE_COUNT=$(grep -c '^MODULE .* exit=0 ' "$WORK_DIR/sched.out" || true)
ev ""
ev "--- scheduler output (one line per compiled module) ---"
cat "$WORK_DIR/sched.out" >> "$EVIDENCE"
ev "--- end scheduler output ---"
log "check exit status: $CHECK_EXIT   elapsed: ${ELAPSED}s"
ev ""
TLOG="$LOGS/$MODULE.log"
if [[ -f "$TLOG" ]]; then
  ev "--- Lean output for $MODFILE (contains the source's own #print axioms) ---"
  cat "$TLOG" >> "$EVIDENCE"; ev "--- end ---"; ev ""
fi
if [[ -f "$LOGS/audit.log" ]]; then
  ev "--- audit file LeanProofUnitAudit.lean ---"; cat "$AUDIT" >> "$EVIDENCE"; ev "--- end ---"; ev ""
  ev "--- Lean output for audit file ---"; cat "$LOGS/audit.log" >> "$EVIDENCE"; ev "--- end ---"; ev ""
fi
# any non-empty output from dependency modules (warnings, sorry, etc.)
for f in "$LOGS"/*.log; do
  [[ "$f" == "$TLOG" || "$f" == "$LOGS/audit.log" ]] && continue
  if [[ -s "$f" ]]; then ev "--- non-empty Lean output: $(basename "$f") ---"; cat "$f" >> "$EVIDENCE"; ev "--- end ---"; fi
done

# ---------------------------------------------------------------- 5. verdict
if (( CHECK_EXIT == 124 || CHECK_EXIT == 137 )); then
  log "RESULT: INCONCLUSIVE (time budget of ${CHECK_TIMEOUT}s exhausted; not a pass)"
  write_receipt "INCONCLUSIVE" 2 "proof-checking invocation exceeded ${CHECK_TIMEOUT}s"; exit 2
fi
# single-line machine report from audit step (c); `#print axioms` text is kept in the evidence
AXLINE=$(grep -h "^AUDIT_AXIOMS: " "$LOGS/audit.log" 2>/dev/null | head -1)
AXIOMS_JSON=$(printf '%s' "$AXLINE" | python3 -I -c 'import sys,json,re; s=sys.stdin.read(); m=re.search(r"\[(.*)\]",s); print(json.dumps(sorted(x.strip() for x in m.group(1).split(",") if x.strip())) if m else "null")')
ok=1
(( CHECK_EXIT == 0 )) || ok=0
grep -q "AUDIT OK:" "$LOGS/audit.log" 2>/dev/null || ok=0
[[ "$AXIOMS_JSON" != "null" ]] || ok=0
for a in $(printf '%s' "$AXIOMS_JSON" | python3 -I -c 'import sys,json; print(" ".join(json.load(sys.stdin) or []))'); do
  [[ " $ALLOWED_AXIOMS " == *" $a "* ]] || { log "disallowed axiom: $a"; ok=0; }
done
if grep -lE "declaration uses .sorry.|error:" "$LOGS"/*.log 2>/dev/null | grep -q .; then log "sorry warning or error found in Lean output of: $(grep -lE "declaration uses .sorry.|error:" "$LOGS"/*.log | xargs -n1 basename | tr '\n' ' ')"; ok=0; fi
log "axioms reported: $AXIOMS_JSON"
log "finished_utc: $(now)"
if (( ok )); then
  log "RESULT: PASS (helper declaration checked; axioms within allowed set)"
  write_receipt "PASS" 0 "declaration elaborated and kernel-checked by Lean $LEAN_VERSION; axioms within allowed set"; exit 0
fi
log "RESULT: FAIL"
write_receipt "FAIL" 1 "Lean check or axiom audit did not succeed (see evidence.txt)"; exit 1