"""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).

Two outputs:

1. An importable Lean file where every database entry is a named declaration. An
   entry is emitted as a real `theorem ... := by <proof>` when a lean-bench run has
   actually proven it (status "proved"), 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.

2. A search index (JSONL) carrying each entry's declaration name, statement, the
   Mathlib identifiers it mentions, and its status -- what search_database.py ranks
   against the current goals at each step of a proof.

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 statement_identifiers(statement):
    """Identifiers mentioned in a statement, for the search index."""
    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('')

    index, 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)
        result = results.get(entry_id)
        proved = bool(result and result.get('success') and result.get('attempts'))
        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()}'
        doc += (' Proven by lean-bench.' if proved
                else ' ASSERTED from the textbook, not proven here.') + ' -/'
        lines.append(doc)

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

        index.append({
            "name": name,
            "id": entry_id,
            "source_label": entry.get('source_label'),
            "kind": entry.get('kind'),
            "status": "proved" if proved else "asserted",
            "statement": entry['statement'],
            "natural_language": entry.get('natural_language', '').split('[fix note:')[0].strip(),
            "identifiers": statement_identifiers(entry['statement']),
            "requires": entry.get('requires', []),
        })
        emitted.append(entry_id)

    return '\n'.join(lines), index, 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('--out-index', type=Path, required=True, help='JSONL search index for search_database.py')
    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, index, 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')
    args.out_index.parent.mkdir(parents=True, exist_ok=True)
    with args.out_index.open('w', encoding='utf-8') as stream:
        for record in index:
            stream.write(json.dumps(record, ensure_ascii=False) + '\n')

    proved = sum(1 for r in index if r['status'] == 'proved')
    print(f'{len(emitted)} premises -> {args.out_lean} ({proved} proven, {len(emitted) - proved} asserted as axioms)')
    print(f'search index -> {args.out_index}')
    if skipped:
        print(f'{len(skipped)} skipped for not being elaboration-checked: {", ".join(skipped)}')


if __name__ == '__main__':
    main()