"""Elaboration-only sanity check for candidate Lean statements from formalize.py.

Checks whether `theorem _ : statement := by sorry` type-checks under the task's
prelude -- i.e. whether the statement is well-formed Lean using identifiers that
actually exist. This does NOT verify the statement is provable, or even that it
faithfully represents the book (that's still on you to spot-check); it only catches
hallucinated Mathlib names and syntax errors for free, before spending any lean-agent
proof attempts on them.

Must run on a machine with `lean`/`lake` installed and, for anything using Mathlib
(almost everything here), a Lake project that already has Mathlib built -- point
--lean-project at it and pass --lean-command '["lake","env","lean"]'. This script
cannot run inside a Cowork sandbox: see lean-bench/textbook/README.md.
"""
import argparse
import json
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path


@dataclass
class ElabResult:
    ok: bool
    diagnostics: str


def elaborates(prelude, statement, command=None, cwd=None, timeout=30):
    command = command or ['lean']
    source = (prelude + '\n' if prelude else '') + 'theorem lean_bench_sanity_check : ' + statement + ' := by sorry\n'
    with tempfile.TemporaryDirectory(prefix='lean-sanity-', dir=cwd) as directory:
        path = Path(directory) / 'SanityCheck.lean'
        path.write_text(source, encoding='utf-8')
        try:
            result = subprocess.run([*command, str(path)], cwd=cwd, capture_output=True,
                                     text=True, timeout=timeout)
        except subprocess.TimeoutExpired:
            return ElabResult(False, 'Lean timeout')
    diagnostics = result.stdout + result.stderr
    ok = result.returncode == 0 and 'error:' not in diagnostics
    return ElabResult(ok, diagnostics[-8000:])


def main():
    parser = argparse.ArgumentParser(description='Check that candidate statements elaborate (not that they are provable)')
    parser.add_argument('dataset', type=Path, help='JSONL with id, statement, optional prelude')
    parser.add_argument('--lean-command', default='["lean"]', help='JSON argv array; e.g. ["lake","env","lean"]')
    parser.add_argument('--lean-project', help='Lean/Lake project working directory (needs Mathlib built)')
    parser.add_argument('--timeout', type=float, default=60)
    parser.add_argument('--report', type=Path, help='Write JSONL of {id, ok, diagnostics} here')
    args = parser.parse_args()

    command = json.loads(args.lean_command)
    tasks = [json.loads(line) for line in args.dataset.read_text(encoding='utf-8').splitlines() if line.strip()]
    reports = []
    failed = 0
    for task in tasks:
        result = elaborates(task.get('prelude', ''), task['statement'], command, args.lean_project, args.timeout)
        reports.append({"id": task['id'], "ok": result.ok, "diagnostics": result.diagnostics})
        print(('OK  ' if result.ok else 'FAIL') + ' ' + task['id'])
        if not result.ok:
            failed += 1

    if args.report:
        with args.report.open('w', encoding='utf-8') as stream:
            for r in reports:
                stream.write(json.dumps(r, ensure_ascii=False) + '\n')

    print(f'{len(tasks) - failed}/{len(tasks)} elaborate cleanly')
    raise SystemExit(1 if failed else 0)


if __name__ == '__main__':
    main()