smoke_signal/scripts/07_codex_exporter.py

#!/usr/bin/env python3
"""
Smoke Signal — Stage 9: Codex Exporter
========================================
Takes approved review decisions and exports clean, schema-valid
JSONL/CSV/Markdown packages ready for Codex ingestion.

Safety rules (non-negotiable):
  - NEVER export unreviewed low-confidence text
  - NEVER mix rights classes in the same export batch
  - NEVER export quarantined or rejected pages
  - ALWAYS link every exported line back to source PDF hash + page
  - ALWAYS validate against schema before writing export

Output formats:
  - codex_export.jsonl  — primary machine-readable format
  - codex_export.csv    — human-readable review copy
  - codex_proof.md      — Markdown proof pack for QA

Usage:
    python scripts/07_codex_exporter.py
    python scripts/07_codex_exporter.py --book-id SS-BOOK-0001
    python scripts/07_codex_exporter.py --batch-id SS-BATCH-001
    python scripts/07_codex_exporter.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

# ── Paths ──────────────────────────────────────────────────────────────────────
ROOT          = Path(__file__).resolve().parents[1]
MANIFEST_CSV  = ROOT / "manifest" / "source_manifest.csv"
REGIONS_DIR   = ROOT / "regions"
CLEANED_DIR   = ROOT / "regions" / "cleaned"
REVIEW_DIR    = ROOT / "review"
EXPORTS_DIR   = ROOT / "exports"
LOGS_DIR      = ROOT / "logs"
SCHEMAS_DIR   = ROOT / "schemas"

EXPORTS_DIR.mkdir(parents=True, exist_ok=True)
LOGS_DIR.mkdir(parents=True, exist_ok=True)

DECISIONS_CSV = REVIEW_DIR / "review_decisions.csv"

# ── Config ────────────────────────────────────────────────────────────────────
CONFIG = {
    "config_version":       "ss_export_v0.1",
    "schema_version":       "1.0",
    "approved_statuses":    ["accepted", "edited"],
    "excluded_statuses":    ["rejected", "quarantined", "illustration-only"],
    "excluded_region_classes": [
        "page-number", "copyright-legal", "publisher-imprint", "decorative-uncertain"
    ],
    "story_region_classes": [
        "narration", "dialogue-speech-bubble", "caption",
        "title", "subtitle", "sign-label"
    ],
    "eligible_rights":      ["public-domain", "licensed-owned", "controlled-internal"],
}

EXPORT_FIELDS = [
    "book_id", "source_id", "source_hash", "page_number", "region_id",
    "page_class", "region_class", "text_raw", "text_clean", "text_final",
    "uncertain_words", "confidence", "review_status", "exclusion_reason",
    "extraction_method", "config_version", "schema_version", "run_id",
]


# ── Manifest I/O ───────────────────────────────────────────────────────────────
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)


# ── Load review decisions ─────────────────────────────────────────────────────
def load_decisions() -> dict:
    decisions = {}
    if not DECISIONS_CSV.exists():
        return decisions
    with open(DECISIONS_CSV, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            decisions[row.get("region_id", "")] = row
    return decisions


# ── Load region data ───────────────────────────────────────────────────────────
def load_regions(book_id: str) -> Optional[dict]:
    """Load cleaned regions if available, otherwise ordered regions."""
    cleaned = CLEANED_DIR / f"{book_id}_cleaned_regions.json"
    if cleaned.exists():
        return json.load(open(cleaned))
    ordered = REGIONS_DIR / f"{book_id}_ordered_regions.json"
    if ordered.exists():
        return json.load(open(ordered))
    return None


# ── Schema validator ──────────────────────────────────────────────────────────
def validate_record(record: dict) -> list:
    """Basic validation — returns list of errors."""
    errors = []
    required = ["book_id", "source_hash", "page_number", "region_id",
                 "region_class", "text_final", "review_status", "run_id"]
    for field in required:
        if not record.get(field):
            errors.append(f"missing_required_field: {field}")

    if record.get("confidence") is not None:
        try:
            conf = float(record["confidence"])
            if not (0.0 <= conf <= 1.0):
                errors.append(f"confidence_out_of_range: {conf}")
        except (ValueError, TypeError):
            errors.append("confidence_not_numeric")

    valid_statuses = CONFIG["approved_statuses"] + ["excluded"]
    if record.get("review_status") not in valid_statuses:
        errors.append(f"invalid_review_status: {record.get('review_status')}")

    return errors


# ── Build export record ────────────────────────────────────────────────────────
def build_export_record(
    region: dict,
    page_data: dict,
    manifest_record: dict,
    decision: Optional[dict],
    run_id: str,
) -> Optional[dict]:
    """
    Build a single Codex export record from region + decision data.
    Returns None if the record should be excluded.
    """
    region_class = region.get("region_class", "decorative-uncertain")
    region_id    = region.get("region_id", "")
    page_num     = region.get("page_number") or page_data.get("page_number")

    # Exclude non-story regions
    if region_class in CONFIG["excluded_region_classes"]:
        return None

    # Check decision status
    if decision:
        review_status = decision.get("status", "pending")
        if review_status not in CONFIG["approved_statuses"]:
            return None
        text_final = decision.get("final_text") or region.get("text_clean") or region.get("text", "")
    else:
        # No decision yet — exclude from export
        return None

    # Determine extraction method
    if region.get("llm_model"):
        method = "ocr+llm"
    elif region.get("text_clean") and region.get("text_clean") != region.get("text"):
        method = "ocr+llm"
    else:
        method = "ocr"

    if page_data.get("route") == "embedded_text":
        method = "embedded-text"

    record = {
        "book_id":          manifest_record.get("book_id", ""),
        "source_id":        manifest_record.get("source_id", ""),
        "source_hash":      manifest_record.get("sha256", ""),
        "page_number":      page_num,
        "region_id":        region_id,
        "page_class":       page_data.get("page_class", "story-page"),
        "region_class":     region_class,
        "text_raw":         region.get("text", ""),
        "text_clean":       region.get("text_clean") or region.get("text", ""),
        "text_final":       text_final.strip(),
        "uncertain_words":  json.dumps(region.get("uncertain_words", [])),
        "confidence":       round(float(region.get("confidence", 0)), 4),
        "review_status":    review_status,
        "exclusion_reason": "",
        "extraction_method": method,
        "config_version":   CONFIG["config_version"],
        "schema_version":   CONFIG["schema_version"],
        "run_id":           run_id,
    }

    return record


# ── Per-book export ───────────────────────────────────────────────────────────
def export_book(
    manifest_record: dict,
    decisions: dict,
    run_id: str,
    dry_run: bool,
) -> tuple:
    book_id     = manifest_record["book_id"]
    filename    = manifest_record["filename"]
    rights      = manifest_record.get("rights_class", "unknown")
    source_hash = manifest_record.get("sha256", "")

    print(f"\n  [{book_id}] {filename}")

    # Rights check
    if rights not in CONFIG["eligible_rights"]:
        print(f"    ✗ Skipped — rights_class={rights} not eligible for export")
        return [], {"error": "rights_not_eligible"}

    # Load region data
    regions_data = load_regions(book_id)
    if regions_data is None:
        print(f"    ✗ No region data found. Run earlier stages first.")
        return [], {"error": "no_region_data"}

    export_records = []
    skipped        = 0
    validation_errors = []

    for page_data in regions_data.get("pages", []):
        page_num = page_data.get("page_number")
        route    = page_data.get("route", "ocr")

        for region in page_data.get("regions", []):
            region_id = region.get("region_id", "")
            decision  = decisions.get(region_id)

            record = build_export_record(
                region, page_data, manifest_record, decision, run_id
            )

            if record is None:
                skipped += 1
                continue

            # Validate
            errors = validate_record(record)
            if errors:
                validation_errors.append({
                    "region_id": region_id,
                    "errors": errors
                })
                print(f"    ⚠️  Validation failed for {region_id}: {errors}")
                continue

            export_records.append(record)

    print(f"    Exported: {len(export_records)} records | Skipped: {skipped} | Validation errors: {len(validation_errors)}")

    if validation_errors:
        print(f"    ⚠️  {len(validation_errors)} records failed validation — excluded from export")

    return export_records, {}


# ── Write export files ────────────────────────────────────────────────────────
def write_exports(
    all_records: list,
    run_id: str,
    batch_id: str,
    dry_run: bool,
) -> dict:
    if not all_records:
        print("\n  No records to export.")
        return {}

    ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
    batch_tag = batch_id or ts

    # ── JSONL ────────────────────────────────────────────────────────────────
    jsonl_path = EXPORTS_DIR / f"codex_export_{batch_tag}.jsonl"
    if not dry_run:
        with open(jsonl_path, "w", encoding="utf-8") as f:
            for rec in all_records:
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")
        print(f"\n  JSONL → {jsonl_path.relative_to(ROOT)}")

    # ── CSV ──────────────────────────────────────────────────────────────────
    csv_path = EXPORTS_DIR / f"codex_export_{batch_tag}.csv"
    if not dry_run:
        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=EXPORT_FIELDS)
            writer.writeheader()
            writer.writerows(all_records)
        print(f"  CSV   → {csv_path.relative_to(ROOT)}")

    # ── Markdown proof pack ──────────────────────────────────────────────────
    md_path = EXPORTS_DIR / f"codex_proof_{batch_tag}.md"
    if not dry_run:
        # Group by book
        by_book = {}
        for rec in all_records:
            bid = rec["book_id"]
            by_book.setdefault(bid, []).append(rec)

        lines = [
            f"# Smoke Signal — Codex Export Proof Pack",
            f"**Batch:** {batch_tag}  ",
            f"**Run ID:** {run_id}  ",
            f"**Exported:** {datetime.utcnow().isoformat()}Z  ",
            f"**Records:** {len(all_records)}  ",
            f"**Config:** {CONFIG['config_version']}  ",
            f"\n---\n",
        ]

        for bid, records in sorted(by_book.items()):
            lines.append(f"## {bid}")
            # Group by page
            by_page = {}
            for rec in records:
                p = rec["page_number"]
                by_page.setdefault(p, []).append(rec)

            for page_num in sorted(by_page.keys()):
                lines.append(f"\n### Page {page_num}")
                for rec in by_page[page_num]:
                    status_icon = "✓" if rec["review_status"] == "accepted" else "✎"
                    lines.append(f"\n**[{status_icon} {rec['region_class']}]** `conf:{rec['confidence']:.0%}`")
                    lines.append(f"\n> {rec['text_final']}\n")
                    if rec.get("uncertain_words") and rec["uncertain_words"] != "[]":
                        lines.append(f"*Uncertain words: {rec['uncertain_words']}*\n")
            lines.append("\n---\n")

        with open(md_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))
        print(f"  MD    → {md_path.relative_to(ROOT)}")

    # ── Export manifest ──────────────────────────────────────────────────────
    manifest_path = EXPORTS_DIR / f"export_manifest_{batch_tag}.json"
    export_manifest = {
        "run_id":        run_id,
        "batch_id":      batch_tag,
        "exported_at":   datetime.utcnow().isoformat() + "Z",
        "config_version": CONFIG["config_version"],
        "schema_version": CONFIG["schema_version"],
        "record_count":  len(all_records),
        "books":         list({r["book_id"] for r in all_records}),
        "files": {
            "jsonl": str(jsonl_path.relative_to(ROOT)) if not dry_run else None,
            "csv":   str(csv_path.relative_to(ROOT)) if not dry_run else None,
            "proof": str(md_path.relative_to(ROOT)) if not dry_run else None,
        }
    }

    if not dry_run:
        with open(manifest_path, "w") as f:
            json.dump(export_manifest, f, indent=2)
        print(f"  Manifest → {manifest_path.relative_to(ROOT)}")

    return export_manifest


# ── Main ───────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Smoke Signal — Stage 9: Codex Exporter")
    parser.add_argument("--book-id",  help="Export a single book by ID")
    parser.add_argument("--batch-id", help="Tag this export batch (e.g. SS-BATCH-001)")
    parser.add_argument("--dry-run",  action="store_true")
    args = parser.parse_args()

    run_id  = f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
    dry_run = args.dry_run

    print(f"\n{'='*60}")
    print(f"  Smoke Signal — Stage 9: Codex Exporter")
    print(f"  Run ID    : {run_id}")
    print(f"  Config    : {CONFIG['config_version']}")
    print(f"  Schema    : {CONFIG['schema_version']}")
    if dry_run:
        print(f"  Mode      : DRY RUN")
    print(f"{'='*60}")

    manifest  = load_manifest()
    decisions = load_decisions()

    if not manifest:
        print("\n  Manifest empty. Run earlier stages first.")
        sys.exit(1)

    if not decisions:
        print("\n  No review decisions found.")
        print("  Run the review workbench (06_review_workbench.py) first.")
        sys.exit(1)

    print(f"\n  Review decisions loaded: {len(decisions)}")
    approved = sum(1 for d in decisions.values() if d.get("status") in CONFIG["approved_statuses"])
    print(f"  Approved decisions: {approved}")

    # Select books
    if args.book_id:
        books = [manifest[args.book_id]] if args.book_id in manifest else []
    else:
        # Export books that have region data
        books = []
        for bid, rec in manifest.items():
            if load_regions(bid) is not None:
                if rec.get("rights_class") in CONFIG["eligible_rights"]:
                    books.append(rec)

    if not books:
        print("\n  No eligible books found for export.")
        sys.exit(0)

    print(f"  Books to export: {len(books)}")

    all_records = []
    t_start     = time.time()

    for record in books:
        records, error = export_book(record, decisions, run_id, dry_run)
        all_records.extend(records)

    # Write all exports
    export_manifest = write_exports(all_records, run_id, args.batch_id, dry_run)

    # Update manifest status to exported
    if not dry_run and all_records:
        exported_books = {r["book_id"] for r in all_records}
        for bid in exported_books:
            if bid in manifest:
                manifest[bid]["status"] = "exported"
        save_manifest(manifest)

    elapsed = round(time.time() - t_start, 1)

    print(f"\n{'─'*60}")
    print(f"  Total records exported : {len(all_records)}")
    print(f"  Time                   : {elapsed}s")
    print(f"{'─'*60}")
    print(f"\n  ✓ Export complete.")
    print(f"  Next: Run validation report (Stage 10) before Codex ingestion.\n")


if __name__ == "__main__":
    main()