smoke_signal/scripts/01_register_sources.py

#!/usr/bin/env python3
"""
Smoke Signal — Stage 1: Source Registry
========================================
Scans /source_pdfs, hashes every PDF, and writes/updates source_manifest.csv.

Usage:
    python scripts/01_register_sources.py
    python scripts/01_register_sources.py --source-dir /path/to/pdfs
    python scripts/01_register_sources.py --validate-only
"""

import hashlib
import csv
import json
import os
import sys
import argparse
from datetime import date
from pathlib import Path
from typing import Optional

try:
    from rich.console import Console
    from rich.table import Table
    from rich.progress import track
    RICH = True
except ImportError:
    RICH = False

# ── Paths ──────────────────────────────────────────────────────────────────────
ROOT         = Path(__file__).resolve().parents[1]
SOURCE_DIR   = ROOT / "source_pdfs"
MANIFEST_CSV = ROOT / "manifest" / "source_manifest.csv"
SCHEMA_FILE  = ROOT / "schemas" / "schema_book_manifest_v1.json"

console = Console() if RICH else None

MANIFEST_FIELDS = [
    "book_id", "source_id", "filename", "sha256", "file_size_bytes",
    "page_count", "rights_class", "source_location", "acquisition_date",
    "status", "allowed_use", "notes"
]

RIGHTS_CLASSES = {"public-domain", "licensed-owned", "controlled-internal", "unknown", "excluded"}


# ── Hashing ────────────────────────────────────────────────────────────────────
def sha256_file(path: Path, chunk: int = 1 << 20) -> str:
    """Return hex SHA-256 of a file, reading in chunks."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(chunk), b""):
            h.update(block)
    return h.hexdigest()


# ── Page count (best-effort, no hard dependency on PDF libs) ──────────────────
def get_page_count(path: Path) -> Optional[int]:
    """Try to count pages without a hard crash if libs are missing."""
    try:
        import fitz  # pymupdf
        doc = fitz.open(str(path))
        count = doc.page_count
        doc.close()
        return count
    except Exception:
        pass
    try:
        import pdfplumber
        with pdfplumber.open(str(path)) as pdf:
            return len(pdf.pages)
    except Exception:
        return None


# ── Book ID generation ─────────────────────────────────────────────────────────
def next_book_id(existing: dict) -> str:
    """Generate the next SS-BOOK-NNNN ID not already in the manifest."""
    used = {v["book_id"] for v in existing.values()}
    for i in range(1, 10_000):
        candidate = f"SS-BOOK-{i:04d}"
        if candidate not in used:
            return candidate
    raise RuntimeError("Ran out of book IDs — this shouldn't happen.")


# ── Manifest I/O ───────────────────────────────────────────────────────────────
def load_manifest() -> dict:
    """Load manifest CSV keyed by sha256."""
    records = {}
    if not MANIFEST_CSV.exists():
        return records
    with open(MANIFEST_CSV, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            if row.get("sha256"):
                records[row["sha256"]] = row
    return records


def save_manifest(records: dict) -> None:
    """Write manifest back to CSV, sorted by book_id."""
    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=MANIFEST_FIELDS)
        writer.writeheader()
        writer.writerows(rows)


# ── Main ───────────────────────────────────────────────────────────────────────
def register_sources(source_dir: Path, validate_only: bool = False) -> None:
    pdf_files = sorted(source_dir.glob("*.pdf"))

    if not pdf_files:
        print(f"[warn] No PDFs found in {source_dir}")
        print("       Add your source PDFs to source_pdfs/ and re-run.")
        return

    print(f"\n{'='*60}")
    print(f"  Smoke Signal — Stage 1: Source Registry")
    print(f"{'='*60}")
    print(f"  Found {len(pdf_files)} PDF(s) in {source_dir}\n")

    existing = load_manifest()
    new_count = 0
    updated_count = 0
    duplicate_count = 0

    iterable = track(pdf_files, description="Hashing PDFs...") if RICH else pdf_files

    for pdf_path in iterable:
        if not RICH:
            print(f"  Processing: {pdf_path.name}")

        file_hash = sha256_file(pdf_path)
        file_size = pdf_path.stat().st_size

        if file_hash in existing:
            rec = existing[file_hash]
            # Update file size if it changed (shouldn't, but track it)
            if rec["filename"] != pdf_path.name:
                print(f"  [dup] {pdf_path.name} → same content as {rec['filename']} ({rec['book_id']})")
                duplicate_count += 1
            continue

        # New source
        page_count = get_page_count(pdf_path)
        book_id    = next_book_id(existing)

        record = {
            "book_id":          book_id,
            "source_id":        "",
            "filename":         pdf_path.name,
            "sha256":           file_hash,
            "file_size_bytes":  file_size,
            "page_count":       page_count if page_count is not None else "",
            "rights_class":     "unknown",        # MUST be set manually
            "source_location":  str(source_dir),
            "acquisition_date": date.today().isoformat(),
            "status":           "pending",
            "allowed_use":      "",
            "notes":            ""
        }
        existing[file_hash] = record
        new_count += 1

        status_icon = "✓" if page_count else "?"
        print(f"  [{status_icon}] Registered {book_id} — {pdf_path.name} ({page_count or '?'} pages)")

    if not validate_only:
        save_manifest(existing)
        print(f"\n  Manifest saved → {MANIFEST_CSV}")

    # ── Summary ───────────────────────────────────────────────────────────────
    total = len(existing)
    unknown_rights = sum(1 for r in existing.values() if r["rights_class"] == "unknown")

    print(f"\n{'─'*60}")
    print(f"  Total sources registered : {total}")
    print(f"  New this run             : {new_count}")
    print(f"  Duplicates skipped       : {duplicate_count}")
    print(f"  Rights class = unknown   : {unknown_rights}  ← ACTION REQUIRED")
    print(f"{'─'*60}\n")

    if unknown_rights > 0:
        print("  ⚠️  ACTION: Open manifest/source_manifest.csv and set")
        print("     rights_class for each book before extraction.")
        print("     Valid values: public-domain | licensed-owned |")
        print("                   controlled-internal | unknown | excluded\n")

    # ── Rights breakdown table ────────────────────────────────────────────────
    rights_counts = {}
    for r in existing.values():
        rc = r.get("rights_class", "unknown")
        rights_counts[rc] = rights_counts.get(rc, 0) + 1

    print("  Rights class breakdown:")
    for rc, count in sorted(rights_counts.items()):
        flag = "  ← EXCLUDED FROM PROCESSING" if rc in ("unknown", "excluded") else ""
        print(f"    {rc:<25} {count}{flag}")
    print()


def validate_manifest() -> bool:
    """Validate all manifest records against JSON schema."""
    try:
        import jsonschema
    except ImportError:
        print("[skip] jsonschema not installed — skipping schema validation")
        return True

    with open(SCHEMA_FILE) as f:
        schema = json.load(f)

    records = load_manifest()
    errors = []
    for sha, rec in records.items():
        # Convert numeric strings for validation
        test_rec = dict(rec)
        if test_rec.get("file_size_bytes"):
            test_rec["file_size_bytes"] = int(test_rec["file_size_bytes"])
        if test_rec.get("page_count"):
            try:
                test_rec["page_count"] = int(test_rec["page_count"])
            except (ValueError, TypeError):
                test_rec["page_count"] = None

        try:
            jsonschema.validate(test_rec, schema)
        except jsonschema.ValidationError as e:
            errors.append(f"  {rec.get('book_id', sha[:8])}: {e.message}")

    if errors:
        print(f"\n  ✗ Schema validation FAILED ({len(errors)} errors):")
        for err in errors:
            print(err)
        return False
    else:
        print(f"\n  ✓ All {len(records)} manifest records pass schema validation.")
        return True


# ── CLI ────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Smoke Signal — Stage 1: Source Registry")
    parser.add_argument("--source-dir",    type=Path, default=SOURCE_DIR,
                        help="Directory containing source PDFs")
    parser.add_argument("--validate-only", action="store_true",
                        help="Only validate existing manifest, don't scan for new files")
    args = parser.parse_args()

    if args.validate_only:
        ok = validate_manifest()
        sys.exit(0 if ok else 1)
    else:
        register_sources(args.source_dir)
        validate_manifest()