"""Verify the database's proofs by compiling them, and record the verdict.

check_statements.py asks "is this statement well-formed Lean"; this asks the much
stronger question "does this proof actually prove it". A database entry is only
allowed to claim proven status after passing here, so nothing enters the citable
premise library on an author's confidence alone.

What counts as passing:
  - `theorem <name> : <statement> := by <proof>` compiles with no error, and
  - the proof text contains no `sorry`/`admit`/`native_decide`, and
  - `#print axioms <name>` reports no axiom dependency beyond Lean's own three
    (propext, Classical.choice, Quot.sound) -- in particular it must not lean on
    `sorryAx`, nor on any of the book-asserted `axiom` declarations, since a
    premise proven *through* an asserted axiom is not independently proven.

With --update, entries that pass get `proof_checked: true` written back into the
database JSONL, which is what makes build_database.py emit them as real theorems
instead of axioms. Needs Lean/Lake and a built Mathlib, same as check_statements.py.
"""
import argparse
import json
import re
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path

BANNED = re.compile(r'\b(sorry|admit|native_decide)\b')
# The three axioms Lean's own standard library is built on; anything else is ours.
ALLOWED_AXIOMS = {'propext', 'Classical.choice', 'Quot.sound'}
AXIOM_LINE = re.compile(r"^'[^']+' depends on axioms: \[(.*)\]$", re.M)


@dataclass
class ProofResult:
    ok: bool
    diagnostics: str


def check_proof(name, statement, proof, prelude='', command=None, cwd=None, timeout=120):
    command = command or ['lean']
    if not proof or not proof.strip():
        return ProofResult(False, 'no proof recorded')
    if BANNED.search(proof):
        return ProofResult(False, f'proof uses a banned escape hatch: {BANNED.search(proof).group(1)}')

    body = '\n'.join('  ' + line for line in proof.strip().splitlines())
    source = ((prelude + '\n' if prelude else '')
              + f'theorem {name} : {statement} := by\n{body}\n'
              + f'#print axioms {name}\n')

    with tempfile.TemporaryDirectory(prefix='lean-proof-', dir=cwd) as directory:
        path = Path(directory) / 'ProofCheck.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 ProofResult(False, 'Lean timeout')

    diagnostics = result.stdout + result.stderr
    if result.returncode != 0 or 'error:' in diagnostics:
        return ProofResult(False, diagnostics[-8000:])

    match = AXIOM_LINE.search(diagnostics)
    if match:
        used = {a.strip() for a in match.group(1).split(',') if a.strip()}
        extra = used - ALLOWED_AXIOMS
        if extra:
            return ProofResult(False, f'proof depends on non-standard axioms: {sorted(extra)}\n'
                                       + diagnostics[-2000:])
    return ProofResult(True, diagnostics[-2000:])


def main():
    parser = argparse.ArgumentParser(description='Compile and verify the proofs recorded in a database JSONL')
    parser.add_argument('database', type=Path, help='e.g. database/goodman-1.6.jsonl')
    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=120)
    parser.add_argument('--report', type=Path, help='write JSONL of {id, ok, diagnostics} here')
    parser.add_argument('--update', action='store_true',
                        help='write proof_checked back into the database for entries that pass')
    parser.add_argument('--only', nargs='*', help='check only these ids')
    args = parser.parse_args()

    command = json.loads(args.lean_command)
    records = [json.loads(line) for line in args.database.read_text(encoding='utf-8').splitlines() if line.strip()]

    reports, checked, failed, skipped = [], 0, 0, 0
    for record in records:
        if args.only and record['id'] not in args.only:
            continue
        if not record.get('proof'):
            skipped += 1
            continue
        name = record['id'].replace('-', '_')
        result = check_proof(name, record['statement'], record['proof'], record.get('prelude', ''),
                             command, args.lean_project, args.timeout)
        reports.append({"id": record['id'], "ok": result.ok, "diagnostics": result.diagnostics})
        print(('PROVED ' if result.ok else 'FAILED ') + record['id'])
        if result.ok:
            checked += 1
            if args.update:
                record['proof_checked'] = True
        else:
            failed += 1
            if args.update:
                record['proof_checked'] = False

    if args.update:
        with args.database.open('w', encoding='utf-8') as stream:
            for record in records:
                stream.write(json.dumps(record, ensure_ascii=False) + '\n')
        print(f'updated proof_checked in {args.database}')

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

    print(f'{checked} proved, {failed} failed, {skipped} entries have no proof recorded yet')
    raise SystemExit(1 if failed else 0)


if __name__ == '__main__':
    main()