"""Assemble a reusable Lean library file from a completed lean-bench run.

Only theorems that actually solved (results.jsonl success == true) go in; each one's
exact winning proof (its last, successful attempt) is reused verbatim, re-wrapped in the
same `theorem <id> : <statement> := by ...` shape lean-bench's verifier already checked,
so the assembled file carries no new proof risk beyond what that run already confirmed --
this script does not itself call Lean.

Plain `import ...` preludes are deduplicated and hoisted to the top of the file; any
other prelude content (a local definition, say) is kept attached to its own theorem
instead of being collapsed, since it isn't necessarily shared with the rest.

A library assembled this way is only as trustworthy as the run it came from: rerun
check_statements.py / lean-bench yourself before importing it into something else you
care about, and treat --module-name as documentation, not a real Lean module boundary
(this script emits one flat .lean file, not a package).
"""
import argparse
import json
from pathlib import Path


def load_jsonl(path):
    return [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines() if line.strip()]


def build(dataset_path, results_path):
    dataset = {t['id']: t for t in load_jsonl(dataset_path)}
    results = {r['id']: r for r in load_jsonl(results_path)}

    shared_imports = set()
    per_theorem_prelude = {}
    for task_id, task in dataset.items():
        prelude = task.get('prelude', '').strip()
        if not prelude:
            continue
        if all(line.strip().startswith('import ') for line in prelude.splitlines()):
            shared_imports.update(line.strip() for line in prelude.splitlines())
        else:
            per_theorem_prelude[task_id] = prelude

    lines = ['-- Generated by textbook/build_library.py -- do not hand-edit proofs;',
             '-- regenerate from a fresh lean-bench run instead.', '']
    lines.extend(sorted(shared_imports))
    lines.append('')

    solved, skipped = [], []
    for task_id, task in dataset.items():
        result = results.get(task_id)
        if not result or not result.get('success') or not result.get('attempts'):
            skipped.append(task_id)
            continue
        proof = result['attempts'][-1]['proof']
        comment = f'-- {task.get("source_label", task_id)}: {task.get("natural_language", "")}'.rstrip()
        lines.append(comment)
        if task_id in per_theorem_prelude:
            lines.append(per_theorem_prelude[task_id])
        lean_name = task_id.replace('-', '_')
        lines.append(f'theorem {lean_name} : {task["statement"]} := by')
        lines.extend('  ' + line for line in proof.splitlines())
        lines.append('')
        solved.append(task_id)

    return '\n'.join(lines), solved, skipped


def main():
    parser = argparse.ArgumentParser(description='Assemble a Lean library file from a solved lean-bench dataset')
    parser.add_argument('dataset', type=Path)
    parser.add_argument('results', type=Path, help="a lean-bench run's results.jsonl")
    parser.add_argument('--out', type=Path, required=True)
    args = parser.parse_args()

    text, solved, skipped = build(args.dataset, args.results)
    args.out.write_text(text + '\n', encoding='utf-8')
    print(f'{len(solved)} theorems assembled -> {args.out}')
    if skipped:
        print(f'{len(skipped)} skipped (not solved in this run): {", ".join(skipped)}')


if __name__ == '__main__':
    main()