scripts/parse_corpus.py

"""Parse all sources — research papers, textbooks, and web tutorials — into chunks.jsonl.

Source priority:
  1. data/papers/*.pdf     → source_type="paper"    (arxiv IDs as source_id)
  2. data/books/*.pdf      → source_type from CANONICAL_TEXT_SOURCES
  3. data/web/*.txt        → source_type from CANONICAL_TEXT_SOURCES

Output: data/chunks.jsonl  (one chunk per line, ready for embedding)

Usage:
    uv run python scripts/parse_corpus.py
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from rich.console import Console
from rich.table import Table

from researchpath.corpus import CANONICAL_RL_PAPERS, CANONICAL_TEXT_SOURCES
from researchpath.parsing import chunk_to_dict, parse_pdf, parse_text

console = Console()
ROOT = Path(__file__).resolve().parents[1]
PAPERS_DIR = ROOT / "data" / "papers"
BOOKS_DIR = ROOT / "data" / "books"
WEB_DIR = ROOT / "data" / "web"
CHUNKS_PATH = ROOT / "data" / "chunks.jsonl"

# Build lookup maps
_PAPER_BY_ID = {p.arxiv_id: p for p in CANONICAL_RL_PAPERS}
_TEXT_BY_FILE = {s.filename: s for s in CANONICAL_TEXT_SOURCES}


def main() -> int:
    table = Table(title="Parse summary")
    table.add_column("Source ID", style="bold")
    table.add_column("Type")
    table.add_column("Tag / Title")
    table.add_column("Chunks", justify="right")
    table.add_column("Chars", justify="right")

    total_chunks = 0
    total_chars = 0

    with open(CHUNKS_PATH, "w", encoding="utf-8") as out:

        # ── 1. Research papers ────────────────────────────────────────────────
        pdfs = sorted(PAPERS_DIR.glob("*.pdf")) if PAPERS_DIR.exists() else []
        if not pdfs:
            console.print("[yellow]No PDFs in data/papers/ — skipping paper corpus.[/yellow]")
        for pdf in pdfs:
            paper_meta = _PAPER_BY_ID.get(pdf.stem)
            tag = paper_meta.tag if paper_meta else "?"
            n = chars = 0
            for chunk in parse_pdf(pdf, source_type="paper"):
                out.write(json.dumps(chunk_to_dict(chunk), ensure_ascii=False) + "\n")
                n += 1
                chars += chunk.n_chars
            table.add_row(pdf.stem, "paper", tag, str(n), f"{chars:,}")
            total_chunks += n
            total_chars += chars

        # ── 2. Book / course PDFs ─────────────────────────────────────────────
        book_pdfs = sorted(BOOKS_DIR.glob("*.pdf")) if BOOKS_DIR.exists() else []
        for pdf in book_pdfs:
            src = _TEXT_BY_FILE.get(pdf.name)
            if src is None:
                console.print(f"[yellow]  Skipping unknown book PDF: {pdf.name}[/yellow]")
                continue
            n = chars = 0
            for chunk in parse_pdf(pdf, source_id=src.source_id, source_type=src.source_type):
                out.write(json.dumps(chunk_to_dict(chunk), ensure_ascii=False) + "\n")
                n += 1
                chars += chunk.n_chars
            short_title = src.title[:40] + ("…" if len(src.title) > 40 else "")
            table.add_row(src.source_id, src.source_type, short_title, str(n), f"{chars:,}")
            total_chunks += n
            total_chars += chars

        # ── 3. Web / tutorial text files ──────────────────────────────────────
        web_txts = sorted(WEB_DIR.glob("*.txt")) if WEB_DIR.exists() else []
        for txt in web_txts:
            src = _TEXT_BY_FILE.get(txt.name)
            if src is None:
                console.print(f"[yellow]  Skipping unknown web file: {txt.name}[/yellow]")
                continue
            content = txt.read_text(encoding="utf-8")
            n = chars = 0
            for chunk in parse_text(content, source_id=src.source_id, source_type=src.source_type):
                out.write(json.dumps(chunk_to_dict(chunk), ensure_ascii=False) + "\n")
                n += 1
                chars += chunk.n_chars
            short_title = src.title[:40] + ("…" if len(src.title) > 40 else "")
            table.add_row(src.source_id, src.source_type, short_title, str(n), f"{chars:,}")
            total_chunks += n
            total_chars += chars

    if total_chunks == 0:
        console.print("[red]No sources found. Run fetch_corpus.py / fetch_pdfs.py / fetch_web_sources.py first.[/red]")
        return 1

    console.print(table)
    console.print(
        f"\n[bold green]Wrote {total_chunks} chunks ({total_chars:,} chars) "
        f"to {CHUNKS_PATH.relative_to(ROOT)}[/bold green]"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())