scripts/extract_books.py

"""
SUPERSEDED by scripts/build_index_from_crawler.py — kept for reference only.
See README.md for the current build path.

Extract plain text from the local book PDFs (../books) into the same on-disk
shape the web crawler produces (scripts/crawl_sources.py) so the RAG chunker
(build_index.py -> iter_crawled_documents) can consume books and articles
through one uniform path: <out>/<slug>/pages/<hash>.txt + <slug>/state.json.

Each book becomes one "source" directory whose single "page" is the whole
book's extracted text. Extraction uses PyMuPDF (fast, clean) with a pypdf
fallback. Books that yield almost no text are flagged (status="low_yield_scanned")
so they can be OCR'd separately instead of silently polluting the index.

Usage:
    python extract_books.py
    python extract_books.py --books-dir ../books --out ../crawled_books
    python extract_books.py --only "100M-Offers" --min-chars 800
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import time
from pathlib import Path

import fitz  # PyMuPDF

# Books whose extracted text is shorter than this are almost certainly
# image-only / scanned PDFs and need OCR, not silent inclusion.
MIN_TEXT_CHARS = 800

# Category is derivable from the ad-hoc filename prefixes the user applied.
CATEGORY_PATTERNS = [
    (re.compile(r"\bseo\b", re.I), "seo"),
    (re.compile(r"\bads?\b|advertis", re.I), "ads"),
    (re.compile(r"\bsocial\b", re.I), "social"),
    (re.compile(r"\bcontent\b", re.I), "content"),
    (re.compile(r"\banalytics\b", re.I), "analytics"),
]


def slugify(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:80] or "book"


def name_hash(name: str) -> str:
    return hashlib.sha1(name.encode("utf-8")).hexdigest()[:16]


def guess_category(filename: str) -> str:
    # Only trust the leading "digital marketing <cat>-" prefix, not the whole title.
    prefix = filename.split("-", 1)[0]
    for pat, cat in CATEGORY_PATTERNS:
        if pat.search(prefix):
            return cat
    return "general"


def clean_text(text: str) -> str:
    text = text.replace("\x0c", "\n")            # form feeds -> newline
    text = re.sub(r"[ \t]+\n", "\n", text)       # trailing spaces
    text = re.sub(r"\n{3,}", "\n\n", text)       # collapse blank-line runs
    # de-hyphenate words broken across line ends: "market-\ning" -> "marketing"
    text = re.sub(r"(\w)-\n(\w)", r"\1\2", text)
    return text.strip()


def extract_pymupdf(path: Path) -> tuple[str, int]:
    doc = fitz.open(str(path))
    try:
        parts = [page.get_text("text") for page in doc]
        return "\n".join(parts), doc.page_count
    finally:
        doc.close()


def extract_pypdf(path: Path) -> tuple[str, int]:
    from pypdf import PdfReader

    reader = PdfReader(str(path))
    parts = [page.extract_text() or "" for page in reader.pages]
    return "\n".join(parts), len(reader.pages)


def extract_book(path: Path) -> tuple[str, int, str]:
    """Return (text, page_count, engine). Tries PyMuPDF, falls back to pypdf."""
    try:
        text, pages = extract_pymupdf(path)
        if len(text.strip()) >= MIN_TEXT_CHARS:
            return clean_text(text), pages, "pymupdf"
    except Exception as exc:
        print(f"  [pymupdf failed] {path.name}: {exc}")
        text, pages = "", 0
    # Fallback (or PyMuPDF yielded almost nothing)
    try:
        alt_text, alt_pages = extract_pypdf(path)
        if len(alt_text.strip()) > len(text.strip()):
            return clean_text(alt_text), alt_pages or pages, "pypdf"
    except Exception as exc:
        print(f"  [pypdf failed] {path.name}: {exc}")
    return clean_text(text), pages, "pymupdf"


def main():
    parser = argparse.ArgumentParser(description="Extract book PDFs into crawl-style text")
    parser.add_argument("--books-dir", default="../books")
    parser.add_argument("--out", default="../crawled_books")
    parser.add_argument("--only", default=None, help="substring filter on filename")
    parser.add_argument("--min-chars", type=int, default=MIN_TEXT_CHARS)
    args = parser.parse_args()

    books_dir = Path(args.books_dir)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    pdfs = sorted(p for p in books_dir.glob("*.pdf") if p.stat().st_size > 0)
    if args.only:
        pdfs = [p for p in pdfs if args.only.lower() in p.name.lower()]

    manifest = []
    ok = low = failed = 0
    print(f"Extracting {len(pdfs)} book PDF(s) -> {out_dir}\n")

    for path in pdfs:
        slug = slugify(path.stem)
        h = name_hash(path.name)
        text, pages, engine = extract_book(path)
        chars = len(text)

        if chars < args.min_chars:
            status = "low_yield_scanned"
            low += 1
        elif chars == 0:
            status = "error"
            failed += 1
        else:
            status = "ok"
            ok += 1

        source_dir = out_dir / slug
        pages_dir = source_dir / "pages"
        pages_dir.mkdir(parents=True, exist_ok=True)
        (pages_dir / f"{h}.txt").write_text(text, encoding="utf-8")

        state = {
            h: {
                "file": path.name,
                "category": guess_category(path.name),
                "status": status,
                "pages": pages,
                "chars": chars,
                "engine": engine,
                "extracted_at": time.time(),
            }
        }
        (source_dir / "state.json").write_text(json.dumps(state, indent=2), encoding="utf-8")

        manifest.append({"slug": slug, "hash": h, **state[h]})
        flag = "" if status == "ok" else f"  <-- {status.upper()}"
        print(f"  [{status:17}] {pages:>4}p {chars:>9,}ch {engine:8} {path.name[:70]}{flag}")

    (out_dir / "extract_manifest.json").write_text(
        json.dumps(manifest, indent=2), encoding="utf-8"
    )
    total_chars = sum(m["chars"] for m in manifest)
    print(
        f"\nDone. {ok} ok, {low} low-yield/scanned, {failed} failed "
        f"| {total_chars:,} chars total across {len(pdfs)} books."
    )
    if low:
        print("Low-yield PDFs (likely need OCR):")
        for m in manifest:
            if m["status"] != "ok":
                print(f"  - {m['file']}  ({m['chars']} chars)")


if __name__ == "__main__":
    main()