"""Generate the citable premise database from a book section's main-content statements.

This is the "third party math library" half of the pipeline: the book's own
Theorem/Proposition/Lemma/Corollary claims, turned into real Lean declarations that
lean-agent can cite by name while proving something else (the bench targets, which
come from the exercises instead -- see formalize_exercises.py).

Output: an importable Lean file where every database entry is a named declaration,
so a proof can cite goodman_1_6_20 by name. An entry is emitted as a real
`theorem ... := by <proof>` once its proof has been verified, and otherwise as an
`axiom` -- asserted from the textbook, which is what the book itself does for its
own body results before an exercise asks you to use them. This is deliberate and
has a real cost: anything proven *through* an axiom inherits it, which is exactly
why the verifier's `#print axioms` check matters. A statement that only elaborates
is not the same as a statement that is true, and a garbled extraction asserted as
an axiom can make downstream "proofs" worthless -- so every axiom here needs to have
been read by a human against the book, and `--require-checked` refuses to emit any
entry not marked as having passed check_statements.py.

There is deliberately no generated search index: search_database.py reads this same
source file and derives the declaration name, identifiers and status itself, so
there is only ever one copy of the data and nothing to keep in sync.

Pass a lean-bench run's results.jsonl with --results to upgrade the entries it
proved from axiom to a real theorem carrying the exact proof that run verified.
"""
import argparse
import json
import re
from pathlib import Path

IDENT_RE = re.compile(r'[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)+|[A-Za-z_][A-Za-z0-9_]*')
# Tokens that say nothing about which premise is relevant.
STOPWORDS = {'let', 'in', 'fun', 'if', 'then', 'else', 'true', 'false', 'Prop', 'Type'}


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


def decl_name(entry_id):
    """goodman-1-6-2-a -> goodman_1_6_2_a (a legal Lean identifier)."""
    return entry_id.replace('-', '_')


def split_prelude(prelude):
    """Return (import lines, other prelude lines) for one entry's prelude."""
    imports, other = [], []
    for line in (prelude or '').splitlines():
        (imports if line.strip().startswith('import ') else other).append(line)
    return imports, '\n'.join(other).strip()


def resolve_proof(entry, result=None):
    """(proof, status) for one entry.

    A proof only counts once Lean has verified it: check_proofs.py sets
    proof_checked on the entry, or a lean-bench run reports success. An unverified
    proof leaves the entry asserted, reported as proof_pending so it is visible
    without being trusted."""
    run_proof = result['attempts'][-1]['proof'] if (
        result and result.get('success') and result.get('attempts')) else None
    own_proof = entry.get('proof') if entry.get('proof_checked') else None
    proof = run_proof or own_proof
    if proof:
        return proof, 'proved'
    if entry.get('proof'):
        return None, 'proof_pending'
    return None, 'asserted'


def statement_identifiers(statement):
    """Identifiers mentioned in a statement, for premise retrieval."""
    found = []
    for match in IDENT_RE.findall(statement or ''):
        if match in STOPWORDS or match in found or len(match) < 2:
            continue
        found.append(match)
    return found


def build(entries, results=None, require_checked=False):
    results = results or {}

    lines = ['-- Generated by textbook/build_database.py -- do not hand-edit.',
             '-- Citable premise database from the textbook\'s main content.',
             '-- `axiom` entries are asserted from the book, NOT proven here: anything',
             '-- proven through one inherits it, so check `#print axioms` downstream.',
             '']

    shared_imports, shared_defs = [], []
    for entry in entries:
        imports, other = split_prelude(entry.get('prelude'))
        for line in imports:
            if line.strip() not in shared_imports:
                shared_imports.append(line.strip())
        # A def can only be declared once per file, so dedupe rather than attaching
        # it per-theorem the way build_library.py does.
        if other and other not in shared_defs:
            shared_defs.append(other)

    lines.extend(sorted(shared_imports))
    lines.append('')
    for block in shared_defs:
        lines.append(block)
        lines.append('')

    emitted, skipped = [], []
    for entry in entries:
        entry_id = entry['id']
        if require_checked and not entry.get('elaborates'):
            skipped.append(entry_id)
            continue

        name = decl_name(entry_id)
        proof, status = resolve_proof(entry, results.get(entry_id))
        proved = status == 'proved'
        proof_pending = status == 'proof_pending'
        provenance = entry.get('source_label', entry_id)

        doc = f'/-- {provenance} (Goodman, Algebra: Abstract and Concrete).'
        if entry.get('natural_language'):
            doc += f' {entry["natural_language"].split("[fix note:")[0].strip()}'
        if proved:
            doc += ' Proof verified by Lean.'
        elif proof_pending:
            doc += ' ASSERTED from the textbook; a proof is recorded but not yet verified.'
        else:
            doc += ' ASSERTED from the textbook, not proven here.'
        doc += ' -/'
        lines.append(doc)

        if proved:
            lines.append(f'theorem {name} : {entry["statement"]} := by')
            lines.extend('  ' + line for line in proof.strip().splitlines())
        else:
            lines.append(f'axiom {name} : {entry["statement"]}')
        lines.append('')

        emitted.append(entry_id)

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


def main():
    parser = argparse.ArgumentParser(description='Generate the citable premise database from main-content statements')
    parser.add_argument('database', type=Path, help='JSONL of main-content statements, e.g. database/goodman-1.6.jsonl')
    parser.add_argument('--results', type=Path, help="a lean-bench run's results.jsonl; proven entries become real theorems")
    parser.add_argument('--out-lean', type=Path, required=True, help='importable Lean file of citable premises')
    parser.add_argument('--require-checked', action='store_true',
                        help='refuse entries not marked "elaborates": true (i.e. not yet verified by check_statements.py)')
    args = parser.parse_args()

    entries = load_jsonl(args.database)
    results = {r['id']: r for r in load_jsonl(args.results)} if args.results else None
    text, emitted, skipped = build(entries, results, args.require_checked)

    args.out_lean.parent.mkdir(parents=True, exist_ok=True)
    args.out_lean.write_text(text + '\n', encoding='utf-8')

    by_id = {e['id']: e for e in entries}
    statuses = [resolve_proof(by_id[i], (results or {}).get(i))[1] for i in emitted]
    proved = statuses.count('proved')
    pending = statuses.count('proof_pending')
    print(f'{len(emitted)} premises -> {args.out_lean} ({proved} proven, '
          f'{len(emitted) - proved} asserted as axioms, of which {pending} have an unverified proof)')
    if skipped:
        print(f'{len(skipped)} skipped for not being elaboration-checked: {", ".join(skipped)}')


if __name__ == '__main__':
    main()