smoke_signal/scripts/02_profile_pdfs.py

#!/usr/bin/env python3
"""
Smoke Signal — Stage 3: PDF Profiler & Renderer
=================================================
For every PDF in the source manifest (status=pending or profiled),
this script:
  1. Opens each PDF and inspects every page
  2. Detects whether a page has embedded text, is image-only, or hybrid
  3. Detects rotation, skew hints, two-page spreads, dimensions
  4. Routes each page: embedded_text | ocr | hybrid
  5. Renders OCR-candidate pages to PNG at 300 DPI (450 DPI fallback)
  6. Saves page_profile.json per book
  7. Updates source_manifest.csv status to 'rendered'

Usage:
    python scripts/02_profile_pdfs.py
    python scripts/02_profile_pdfs.py --book-id SS-BOOK-0001
    python scripts/02_profile_pdfs.py --batch-id SS-BATCH-001
    python scripts/02_profile_pdfs.py --dry-run
"""

import argparse
import csv
import json
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

ROOT         = Path(__file__).resolve().parents[1]
SOURCE_DIR   = ROOT / "source_pdfs"
MANIFEST_CSV = ROOT / "manifest" / "source_manifest.csv"
RENDERS_DIR  = ROOT / "renders"
PROFILES_DIR = ROOT / "manifest" / "page_profiles"
LOGS_DIR     = ROOT / "logs"

PROFILES_DIR.mkdir(parents=True, exist_ok=True)
RENDERS_DIR.mkdir(parents=True, exist_ok=True)
LOGS_DIR.mkdir(parents=True, exist_ok=True)

CONFIG = {
    "config_version":        "ss_profiler_v0.1",
    "render_dpi_baseline":   300,
    "render_dpi_fallback":   450,
    "text_char_threshold":   20,
    "spread_aspect_ratio":   1.6,
    "image_coverage_threshold": 0.15,
    "skip_statuses":         ["exported", "quarantined"],
    "eligible_statuses":     ["pending", "profiled"],
}


def load_manifest() -> dict:
    records = {}
    if not MANIFEST_CSV.exists():
        return records
    with open(MANIFEST_CSV, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            if row.get("book_id"):
                records[row["book_id"]] = row
    return records


def save_manifest(records: dict) -> None:
    fields = [
        "book_id", "source_id", "filename", "sha256", "file_size_bytes",
        "page_count", "rights_class", "source_location", "acquisition_date",
        "status", "allowed_use", "notes"
    ]
    rows = sorted(records.values(), key=lambda r: r.get("book_id", ""))
    with open(MANIFEST_CSV, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def _import_fitz():
    try:
        import fitz
        return fitz
    except ImportError:
        print("  [error] PyMuPDF (fitz) not installed. Run: pip install pymupdf")
        sys.exit(1)


def analyse_page(page, fitz, config: dict) -> dict:
    rect   = page.rect
    width  = rect.width
    height = rect.height
    rotation = page.rotation
    text = page.get_text("text").strip()
    char_count = len(text)
    has_embedded_text = char_count >= config["text_char_threshold"]
    image_list = page.get_images(full=True)
    has_images = len(image_list) > 0
    image_coverage = 0.0
    page_area = width * height
    if has_images and page_area > 0:
        for img in image_list:
            try:
                rects = page.get_image_rects(img[0])
                for r in rects:
                    image_coverage += abs(r.width * r.height) / page_area
            except Exception:
                pass
        image_coverage = min(image_coverage, 1.0)
    if has_embedded_text and not has_images:
        route = "embedded_text"
    elif has_embedded_text and has_images:
        route = "hybrid" if image_coverage >= config["image_coverage_threshold"] else "embedded_text"
    else:
        route = "ocr"
    aspect = width / height if height > 0 else 0
    is_spread = aspect >= config["spread_aspect_ratio"]
    warnings = []
    if rotation not in (0, 360):
        warnings.append(f"rotation_{rotation}deg")
    if is_spread:
        warnings.append("possible_two_page_spread")
    if route == "ocr" and char_count == 0 and not has_images:
        warnings.append("blank_or_undetectable_page")
    return {
        "width_pt": round(width, 2), "height_pt": round(height, 2),
        "rotation_deg": rotation, "char_count": char_count,
        "has_embedded_text": has_embedded_text, "has_images": has_images,
        "image_count": len(image_list), "image_coverage": round(image_coverage, 3),
        "is_spread": is_spread, "route": route, "warnings": warnings,
    }


def render_page(page, book_id: str, page_num: int, dpi: int, render_dir: Path) -> Optional[str]:
    book_render_dir = render_dir / book_id
    book_render_dir.mkdir(parents=True, exist_ok=True)
    filename = f"{book_id}_page_{page_num:04d}_{dpi}dpi.png"
    out_path = book_render_dir / filename
    try:
        mat = page.fitz_module.Matrix(dpi / 72, dpi / 72)
        pix = page.get_pixmap(matrix=mat, alpha=False)
        pix.save(str(out_path))
        return str(out_path.relative_to(ROOT))
    except Exception:
        return None


def profile_book(record: dict, dry_run: bool = False):
    book_id  = record["book_id"]
    filename = record["filename"]
    pdf_path = SOURCE_DIR / filename
    print(f"\n  [{book_id}] {filename}")
    if not pdf_path.exists():
        return record, None, {"error": "file_not_found"}
    if record.get("rights_class") in ("unknown", "excluded"):
        return record, None, {"error": "rights_blocked"}
    fitz = _import_fitz()
    try:
        doc = fitz.open(str(pdf_path))
    except Exception as e:
        return record, None, {"error": str(e)}
    page_count = doc.page_count
    print(f"    Pages: {page_count}")
    pages = []
    render_errors = []
    route_counts = {"embedded_text": 0, "ocr": 0, "hybrid": 0}
    for i in range(page_count):
        page_num = i + 1
        page = doc[i]
        page.fitz_module = fitz
        profile = analyse_page(page, fitz, CONFIG)
        profile["page_number"] = page_num
        render_path = None
        render_dpi  = None
        if not dry_run and profile["route"] in ("ocr", "hybrid"):
            dpi = CONFIG["render_dpi_baseline"]
            render_path = render_page(page, book_id, page_num, dpi, RENDERS_DIR)
            if render_path is None:
                dpi = CONFIG["render_dpi_fallback"]
                render_path = render_page(page, book_id, page_num, dpi, RENDERS_DIR)
                if render_path is None:
                    render_errors.append(page_num)
                else:
                    render_dpi = dpi
            else:
                render_dpi = dpi
        profile["render_path"] = render_path
        profile["render_dpi"]  = render_dpi
        pages.append(profile)
        route_counts[profile["route"]] += 1
    doc.close()
    book_profile = {
        "book_id": book_id, "filename": filename,
        "source_hash": record.get("sha256", ""),
        "page_count": page_count, "config_version": CONFIG["config_version"],
        "profiled_at": datetime.utcnow().isoformat() + "Z",
        "route_summary": route_counts, "render_errors": render_errors, "pages": pages,
    }
    if not dry_run:
        profile_path = PROFILES_DIR / f"{book_id}_page_profile.json"
        with open(profile_path, "w", encoding="utf-8") as f:
            json.dump(book_profile, f, indent=2)
        print(f"    Profile saved -> {profile_path.relative_to(ROOT)}")
    total_ocr = route_counts["ocr"] + route_counts["hybrid"]
    print(f"    Routes: embedded={route_counts['embedded_text']} | ocr={route_counts['ocr']} | hybrid={route_counts['hybrid']}")
    record["page_count"] = page_count
    record["status"]     = "profiled" if not dry_run else record["status"]
    return record, book_profile, {}


def main():
    parser = argparse.ArgumentParser(description="Smoke Signal Stage 3: PDF Profiler")
    parser.add_argument("--book-id", help="Profile a single book by ID")
    parser.add_argument("--batch-id", help="Tag this run with a batch ID")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--all", action="store_true")
    args = parser.parse_args()
    run_id  = args.batch_id or f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
    dry_run = args.dry_run
    print(f"\nSmoke Signal Stage 3 | {run_id} | config={CONFIG['config_version']}")
    manifest = load_manifest()
    if not manifest:
        print("Manifest empty. Run 01_register_sources.py first.")
        sys.exit(1)
    eligible = CONFIG["eligible_statuses"] if not args.all else ["pending","profiled","rendered"]
    books = [manifest[args.book_id]] if args.book_id else [r for r in manifest.values() if r.get("status") in eligible]
    if not books:
        print(f"No books with status in {eligible}.")
        sys.exit(0)
    print(f"Books to profile: {len(books)}")
    results = []
    t_start = time.time()
    for record in books:
        updated, profile, error = profile_book(record, dry_run=dry_run)
        result = {"book_id": record["book_id"], "filename": record["filename"], "error": error}
        if profile:
            result.update({"route_summary": profile["route_summary"], "page_count": profile["page_count"]})
            manifest[record["book_id"]] = updated
        results.append(result)
    if not dry_run:
        save_manifest(manifest)
        log_path = LOGS_DIR / f"{run_id}_profile_run.json"
        with open(log_path, "w") as f:
            json.dump({"run_id": run_id, "config": CONFIG, "results": results}, f, indent=2)
    elapsed = round(time.time() - t_start, 1)
    succeeded = sum(1 for r in results if not r.get("error"))
    print(f"\nDone: {succeeded}/{len(results)} succeeded in {elapsed}s")
    print("Next: run 03_ocr_bakeoff.py on calibration corpus")


if __name__ == "__main__":
    main()