"""Slice one book section out of a `pdftotext -layout` dump and list its labeled
Theorem/Proposition/Lemma/Corollary/Definition statements, as input for formalize.py.

Produces two files: the section's raw text, and a JSON checklist of labeled items
(kind, label, first line) in book order. The checklist exists so formalize.py's LLM
step has an explicit inventory to cover and cannot silently skip or invent items.
"""
import argparse
import json
import re
from pathlib import Path

LABEL_RE = re.compile(r'^(Definition|Theorem|Proposition|Lemma|Corollary)\s+(\d+\.\d+(?:\.\d+)?)\.?\s*(.*)$')


def find_heading_line(lines, heading):
    for i, line in enumerate(lines):
        if line.strip() == heading:
            return i
    raise ValueError(f'Heading not found: {heading!r}')


def extract_section(full_text, start_heading, end_heading):
    lines = full_text.split('\n')
    start = find_heading_line(lines, start_heading)
    end = find_heading_line(lines, end_heading) if end_heading else len(lines)
    if end <= start:
        raise ValueError('end_heading must come after start_heading')
    return '\n'.join(lines[start:end])


def labeled_statements(section_text, section_prefix):
    """Unique labeled items in section_text whose label starts with section_prefix
    (e.g. '1.6.'), keeping only the first occurrence of each label: later matches are
    almost always proof continuations across a page break or cross-references, not new
    statements. Returned sorted in book order (numeric, not lexical)."""
    seen = {}
    for line in section_text.split('\n'):
        line = line.strip()
        m = LABEL_RE.match(line)
        if not m:
            continue
        kind, label, rest = m.group(1), m.group(2), m.group(3)
        if not label.startswith(section_prefix) or label in seen:
            continue
        seen[label] = {"kind": kind, "label": label, "first_line": line[:120]}

    def sort_key(item):
        return tuple(int(p) for p in item['label'].split('.'))

    return sorted(seen.values(), key=sort_key)


def main():
    parser = argparse.ArgumentParser(description='Slice one section out of a pdftotext -layout dump')
    parser.add_argument('layout_txt', type=Path)
    parser.add_argument('--start-heading', required=True, help="e.g. '1.6. Divisibility in the Integers'")
    parser.add_argument('--end-heading', help="e.g. '1.7. Modular Arithmetic'; omit for end of file")
    parser.add_argument('--section-prefix', required=True, help="e.g. '1.6.'")
    parser.add_argument('--out-text', type=Path, required=True)
    parser.add_argument('--out-index', type=Path, required=True)
    args = parser.parse_args()

    full_text = args.layout_txt.read_text(encoding='utf-8')
    section_text = extract_section(full_text, args.start_heading, args.end_heading)
    args.out_text.write_text(section_text, encoding='utf-8')
    index = labeled_statements(section_text, args.section_prefix)
    args.out_index.write_text(json.dumps(index, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')
    print(f'{len(index)} labeled statements -> {args.out_index}')


if __name__ == '__main__':
    main()