#!/usr/bin/env python3
"""
Smoke Signal — Stage 4: OCR Bake-Off
======================================
Runs Surya (primary) and Tesseract (fallback) on rendered page images
from the calibration corpus. Compares outputs, measures accuracy against
gold set if available, and freezes a baseline OCR config.

Inputs:
    renders/<BOOK_ID>/<BOOK_ID>_page_NNNN_300dpi.png
    manifest/page_profiles/<BOOK_ID>_page_profile.json
    manifest/source_manifest.csv

Outputs:
    ocr_raw/<BOOK_ID>/<BOOK_ID>_page_NNNN_ocr.json   — per-page OCR result
    manifest/ocr_run_<RUN_ID>.json                    — run summary
    configs/ss_ocr_config_v0.1.json                   — frozen baseline config

Usage:
    python scripts/03_ocr_bakeoff.py
    python scripts/03_ocr_bakeoff.py --book-id SS-BOOK-0001
    python scripts/03_ocr_bakeoff.py --engine surya
    python scripts/03_ocr_bakeoff.py --engine tesseract
    python scripts/03_ocr_bakeoff.py --engine both
    python scripts/03_ocr_bakeoff.py --dry-run
"""

import argparse
import csv
import hashlib
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"
PROFILES_DIR = ROOT / "manifest" / "page_profiles"
RENDERS_DIR  = ROOT / "renders"
OCR_RAW_DIR  = ROOT / "ocr_raw"
CONFIGS_DIR  = ROOT / "configs"
LOGS_DIR     = ROOT / "logs"

OCR_RAW_DIR.mkdir(parents=True, exist_ok=True)
CONFIGS_DIR.mkdir(parents=True, exist_ok=True)
LOGS_DIR.mkdir(parents=True, exist_ok=True)

# ── Frozen OCR config (do not change mid-batch) ────────────────────────────────
OCR_CONFIG = {
    "config_version":      "ss_ocr_config_v0.1",
    "primary_engine":      "surya",
    "fallback_engine":     "tesseract",
    "surya_langs":         ["en"],
    "surya_det_batch":     4,
    "surya_rec_batch":     4,
    "tesseract_lang":      "eng",
    "tesseract_psm":       6,          # assume uniform block of text
    "confidence_threshold_auto_accept": 0.85,
    "confidence_threshold_review":      0.60,
    "confidence_threshold_quarantine":  0.40,
    "eligible_routes":     ["ocr", "hybrid"],
}


# ── Manifest / profile loaders ─────────────────────────────────────────────────
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 load_page_profile(book_id: str) -> Optional[dict]:
    profile_path = PROFILES_DIR / f"{book_id}_page_profile.json"
    if not profile_path.exists():
        return None
    with open(profile_path, encoding="utf-8") as f:
        return json.load(f)


# ── Surya OCR ──────────────────────────────────────────────────────────────────
def _safe_load_surya_component(loader, checkpoint: Optional[str]):
    """
    Attempt to pass checkpoint to Surya loader, with safe fallback for older APIs.
    """
    if not checkpoint:
        return loader()
    try:
        return loader(checkpoint=checkpoint)
    except TypeError:
        return loader()


def load_surya_context(checkpoint: Optional[str] = None) -> Optional[dict]:
    """
    Load Surya OCR models once per run.
    Tries modern predictor API first, then legacy API.
    Returns context dict or None if Surya import/loading fails.
    """
    # New API (surya-ocr>=0.17 style)
    try:
        from surya.foundation import FoundationPredictor
        from surya.detection import DetectionPredictor
        from surya.recognition import RecognitionPredictor
        try:
            from surya.common.surya.schema import TaskNames
            task_name = TaskNames.ocr_with_boxes
        except Exception:
            task_name = "ocr_with_boxes"

        foundation_predictor = _safe_load_surya_component(FoundationPredictor, checkpoint)
        det_predictor = DetectionPredictor()
        rec_predictor = RecognitionPredictor(foundation_predictor)
        return {
            "api": "predictor-v2",
            "task_name": task_name,
            "det_predictor": det_predictor,
            "rec_predictor": rec_predictor,
            "checkpoint": checkpoint,
        }
    except Exception:
        pass

    # Legacy API (surya-ocr<=0.6 style)
    try:
        from surya.ocr import run_ocr
        from surya.model.detection.model import load_model as load_det_model
        from surya.model.detection.processor import load_processor as load_det_processor
        from surya.model.recognition.model import load_model as load_rec_model
        from surya.model.recognition.processor import load_processor as load_rec_processor
    except ImportError:
        return None

    try:
        det_model = _safe_load_surya_component(load_det_model, checkpoint)
        det_processor = _safe_load_surya_component(load_det_processor, checkpoint)
        rec_model = _safe_load_surya_component(load_rec_model, checkpoint)
        rec_processor = _safe_load_surya_component(load_rec_processor, checkpoint)
        return {
            "run": run_ocr,
            "det_model": det_model,
            "det_processor": det_processor,
            "rec_model": rec_model,
            "rec_processor": rec_processor,
            "checkpoint": checkpoint,
        }
    except Exception:
        return None


def _run_surya(image_path: Path, langs: list, surya_ctx: Optional[dict] = None) -> dict:
    """
    Run Surya OCR on a single page image.
    Returns standardised result dict.
    """
    try:
        from PIL import Image
    except ImportError as e:
        return {
            "engine": "surya",
            "error":  f"Import error: {e}. Run: pip install surya-ocr",
            "text":   "",
            "words":  [],
            "confidence": 0.0,
        }

    ctx = surya_ctx or load_surya_context()
    if not ctx:
        return {
            "engine": "surya",
            "error":  "Surya model load failed. Check surya-ocr install and checkpoint path.",
            "text":   "",
            "words":  [],
            "confidence": 0.0,
        }

    try:
        image = Image.open(str(image_path)).convert("RGB")
        if ctx.get("api") == "predictor-v2":
            results = ctx["rec_predictor"](
                [image],
                task_names=[ctx["task_name"]],
                det_predictor=ctx["det_predictor"],
                highres_images=[image],
                math_mode=True,
            )
        else:
            results = ctx["run"](
                [image],
                [langs],
                ctx["det_model"],
                ctx["det_processor"],
                ctx["rec_model"],
                ctx["rec_processor"],
            )

        page_result = results[0]

        # Extract text and confidence from Surya's TextLine objects
        words     = []
        full_text = []
        confidences = []

        for line in page_result.text_lines:
            text = line.text.strip()
            conf = float(line.confidence) if hasattr(line, "confidence") else 1.0
            if text:
                full_text.append(text)
                confidences.append(conf)
                bbox = getattr(line, "bbox", None)
                if bbox is None:
                    bbox = getattr(line, "polygon", None)
                words.append({
                    "text":       text,
                    "confidence": round(conf, 4),
                    "bbox":       bbox,
                })

        avg_conf = round(sum(confidences) / len(confidences), 4) if confidences else 0.0

        return {
            "engine":     "surya",
            "text":       "\n".join(full_text),
            "words":      words,
            "confidence": avg_conf,
            "line_count": len(words),
            "model_checkpoint": ctx.get("checkpoint") or "base",
            "error":      None,
        }

    except Exception as e:
        return {
            "engine":     "surya",
            "error":      str(e),
            "text":       "",
            "words":      [],
            "confidence": 0.0,
            "model_checkpoint": ctx.get("checkpoint") or "base",
        }


# ── Tesseract OCR (fallback) ───────────────────────────────────────────────────
def _run_tesseract(image_path: Path, lang: str = "eng", psm: int = 6) -> dict:
    """
    Run Tesseract on a single page image.
    Returns standardised result dict.
    """
    try:
        import pytesseract
        from PIL import Image
    except ImportError as e:
        return {
            "engine":     "tesseract",
            "error":      f"Import error: {e}. Run: pip install pytesseract pillow",
            "text":       "",
            "words":      [],
            "confidence": 0.0,
        }

    try:
        image = Image.open(str(image_path)).convert("RGB")
        config = f"--psm {psm}"

        # Get word-level data with confidence
        data = pytesseract.image_to_data(
            image,
            lang=lang,
            config=config,
            output_type=pytesseract.Output.DICT,
        )

        words = []
        confidences = []
        full_text_parts = []

        for i, word_text in enumerate(data["text"]):
            word_text = str(word_text).strip()
            conf = int(data["conf"][i])
            if word_text and conf > 0:
                conf_norm = conf / 100.0
                words.append({
                    "text":       word_text,
                    "confidence": round(conf_norm, 4),
                    "bbox": [
                        data["left"][i], data["top"][i],
                        data["left"][i] + data["width"][i],
                        data["top"][i] + data["height"][i],
                    ],
                })
                confidences.append(conf_norm)
                full_text_parts.append(word_text)

        avg_conf = round(sum(confidences) / len(confidences), 4) if confidences else 0.0
        full_text = pytesseract.image_to_string(image, lang=lang, config=config).strip()

        return {
            "engine":     "tesseract",
            "text":       full_text,
            "words":      words,
            "confidence": avg_conf,
            "line_count": len(words),
            "error":      None,
        }

    except Exception as e:
        return {
            "engine":     "tesseract",
            "error":      str(e),
            "text":       "",
            "words":      [],
            "confidence": 0.0,
        }


# ── Confidence gate ────────────────────────────────────────────────────────────
def confidence_gate(confidence: float, config: dict) -> str:
    """Return auto-accept | review-required | quarantine based on thresholds."""
    if confidence >= config["confidence_threshold_auto_accept"]:
        return "auto-accept"
    elif confidence >= config["confidence_threshold_review"]:
        return "review-required"
    elif confidence >= config["confidence_threshold_quarantine"]:
        return "quarantine"
    else:
        return "quarantine"


# ── Per-page OCR ───────────────────────────────────────────────────────────────
def ocr_page(
    image_path: Path,
    book_id: str,
    page_num: int,
    engine: str,
    config: dict,
    surya_ctx: Optional[dict] = None,
    dry_run: bool = False,
) -> dict:
    """Run OCR on one page, save result, return summary."""

    result = {
        "book_id":    book_id,
        "page_number": page_num,
        "image_path": str(image_path),
        "engine":     engine,
        "ocr_at":     datetime.utcnow().isoformat() + "Z",
        "config_version": config["config_version"],
    }

    if dry_run:
        result.update({
            "text": "[dry-run]", "confidence": 0.0,
            "gate": "dry-run", "error": None, "words": [],
        })
        return result

    if engine == "surya":
        ocr_out = _run_surya(image_path, config["surya_langs"], surya_ctx=surya_ctx)
    elif engine == "tesseract":
        ocr_out = _run_tesseract(image_path, config["tesseract_lang"], config["tesseract_psm"])
    else:
        ocr_out = {"engine": engine, "error": f"Unknown engine: {engine}", "text": "", "words": [], "confidence": 0.0}

    result.update(ocr_out)
    result["gate"] = confidence_gate(result.get("confidence", 0.0), config)

    # Save per-page OCR JSON
    book_ocr_dir = OCR_RAW_DIR / book_id
    book_ocr_dir.mkdir(parents=True, exist_ok=True)
    out_path = book_ocr_dir / f"{book_id}_page_{page_num:04d}_{engine}_ocr.json"
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)

    return result


# ── Per-book OCR runner ────────────────────────────────────────────────────────
def ocr_book(
    record: dict,
    engine: str,
    config: dict,
    surya_ctx: Optional[dict] = None,
    dry_run: bool = False,
) -> dict:
    book_id  = record["book_id"]
    print(f"\n  [{book_id}] {record['filename']} — engine: {engine}")

    profile = load_page_profile(book_id)
    if not profile:
        print(f"    ✗ No page profile found. Run 02_profile_pdfs.py first.")
        return {"book_id": book_id, "error": "no_profile", "pages": []}

    eligible_routes = config["eligible_routes"]
    ocr_pages = [p for p in profile["pages"] if p.get("route") in eligible_routes]

    print(f"    OCR-eligible pages: {len(ocr_pages)} / {profile['page_count']}")

    if not ocr_pages:
        print(f"    ✓ No OCR pages — all embedded text.")
        return {"book_id": book_id, "error": None, "pages": [], "skipped": True}

    page_results   = []
    confidences    = []
    gate_counts    = {"auto-accept": 0, "review-required": 0, "quarantine": 0, "dry-run": 0}
    errors         = []

    for page_info in ocr_pages:
        page_num    = page_info["page_number"]
        render_path = page_info.get("render_path")

        if not render_path:
            # Try to find render file
            render_path_candidates = list((RENDERS_DIR / book_id).glob(
                f"{book_id}_page_{page_num:04d}_*.png"
            )) if (RENDERS_DIR / book_id).exists() else []
            render_path = str(render_path_candidates[0]) if render_path_candidates else None

        if not render_path:
            print(f"    ⚠ Page {page_num}: no render found — skipping")
            errors.append({"page": page_num, "error": "no_render"})
            continue

        image_path = Path(render_path) if Path(render_path).is_absolute() else ROOT / render_path

        if not image_path.exists():
            print(f"    ⚠ Page {page_num}: render file missing — {image_path}")
            errors.append({"page": page_num, "error": "render_missing"})
            continue

        result = ocr_page(image_path, book_id, page_num, engine, config, surya_ctx=surya_ctx, dry_run=dry_run)
        page_results.append(result)

        conf = result.get("confidence", 0.0)
        gate = result.get("gate", "quarantine")
        confidences.append(conf)
        gate_counts[gate] = gate_counts.get(gate, 0) + 1

        status = "✓" if gate == "auto-accept" else "⚠" if gate == "review-required" else "✗"
        print(f"    {status} p{page_num:03d} conf={conf:.2f} gate={gate}")

    avg_conf = round(sum(confidences) / len(confidences), 4) if confidences else 0.0

    print(f"    Avg confidence : {avg_conf:.2f}")
    print(f"    Gates          : {gate_counts}")
    if errors:
        print(f"    Errors         : {len(errors)}")

    return {
        "book_id":      book_id,
        "engine":       engine,
        "pages_ocred":  len(page_results),
        "avg_confidence": avg_conf,
        "gate_counts":  gate_counts,
        "errors":       errors,
        "error":        None,
    }


# ── Save frozen config ─────────────────────────────────────────────────────────
def save_ocr_config(config: dict) -> None:
    config_path = CONFIGS_DIR / f"{config['config_version']}.json"
    if not config_path.exists():
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump({
                **config,
                "frozen_at": datetime.utcnow().isoformat() + "Z",
                "note": "DO NOT change this file mid-batch. Create a new version instead.",
            }, f, indent=2)
        print(f"\n  Config frozen → {config_path.relative_to(ROOT)}")
    else:
        print(f"\n  Config already exists → {config_path.relative_to(ROOT)} (not overwritten)")


# ── Main ───────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Smoke Signal — Stage 4: OCR Bake-Off")
    parser.add_argument("--book-id",  help="OCR a single book by ID")
    parser.add_argument("--batch-id", help="Tag this run with a batch ID")
    parser.add_argument("--engine",   choices=["surya", "tesseract", "both"],
                        default="surya", help="OCR engine to use (default: surya)")
    parser.add_argument("--model",    default=None,
                        help="Optional Surya checkpoint/model path for OCR engine=surya")
    parser.add_argument("--dry-run",  action="store_true", help="No files written")
    parser.add_argument("--all",      action="store_true", help="Include already-OCRed books")
    args = parser.parse_args()

    run_id  = args.batch_id or f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
    engines = ["surya", "tesseract"] if args.engine == "both" else [args.engine]
    run_config = dict(OCR_CONFIG)
    if args.model:
        run_config["surya_model_checkpoint"] = args.model
        # Governance: checkpoint changes require a new config version.
        model_tag = hashlib.sha256(args.model.encode("utf-8")).hexdigest()[:8]
        run_config["config_version"] = f"{OCR_CONFIG['config_version']}_ft_{model_tag}"

    print(f"\n{'='*60}")
    print(f"  Smoke Signal — Stage 4: OCR Bake-Off")
    print(f"  Run ID  : {run_id}")
    print(f"  Engines : {engines}")
    print(f"  Config  : {run_config['config_version']}")
    if args.model:
        print(f"  Surya model override : {args.model}")
    if args.dry_run:
        print(f"  Mode    : DRY RUN")
    print(f"{'='*60}")

    manifest = load_manifest()
    if not manifest:
        print("\n  [error] Manifest empty. Run 01_register_sources.py first.")
        sys.exit(1)

    eligible_statuses = ["profiled", "rendered"] if not args.all else \
                        ["profiled", "rendered", "ocred"]

    if args.book_id:
        books = [manifest[args.book_id]] if args.book_id in manifest else []
        if not books:
            print(f"  [error] Book {args.book_id} not in manifest.")
            sys.exit(1)
    else:
        books = [r for r in manifest.values() if r.get("status") in eligible_statuses]

    # Governance: never process unknown/excluded rights in OCR batches.
    books = [r for r in books if r.get("rights_class") not in ("unknown", "excluded")]

    if not books:
        print(f"\n  No books eligible after status/rights filters.")
        print(f"  Eligible statuses: {eligible_statuses}")
        print("  Rights blocked: unknown, excluded")
        sys.exit(0)

    print(f"\n  Books to OCR: {len(books)}")

    all_results = []
    t_start     = time.time()
    surya_ctx = None
    if "surya" in engines and not args.dry_run:
        surya_ctx = load_surya_context(args.model)
        if not surya_ctx:
            print("\n  [error] Could not load Surya models/checkpoint.")
            print("  Check surya-ocr install and --model path.")
            sys.exit(1)

    for record in books:
        for engine in engines:
            result = ocr_book(record, engine, run_config, surya_ctx=surya_ctx, dry_run=args.dry_run)
            all_results.append(result)

            # Update manifest status
            if not result.get("error") and not args.dry_run:
                manifest[record["book_id"]]["status"] = "ocred"

    # Save manifest + config + run log
    if not args.dry_run:
        save_manifest(manifest)
        save_ocr_config(run_config)

        log_path = LOGS_DIR / f"{run_id}_ocr_bakeoff.json"
        with open(log_path, "w", encoding="utf-8") as f:
            json.dump({
                "run_id":   run_id,
                "run_at":   datetime.utcnow().isoformat() + "Z",
                "engines":  engines,
                "config":   run_config,
                "surya_model_checkpoint": args.model or "base",
                "results":  all_results,
            }, f, indent=2)
        print(f"\n  Run log → {log_path.relative_to(ROOT)}")

    # ── Summary ───────────────────────────────────────────────────────────────
    elapsed   = round(time.time() - t_start, 1)
    succeeded = sum(1 for r in all_results if not r.get("error"))
    total_pages = sum(r.get("pages_ocred", 0) for r in all_results)
    avg_conf  = (
        sum(r.get("avg_confidence", 0) for r in all_results if not r.get("error")) / max(succeeded, 1)
    )

    print(f"\n{'─'*60}")
    print(f"  Books processed : {len(books)}")
    print(f"  Runs succeeded  : {succeeded}")
    print(f"  Pages OCR-ed    : {total_pages}")
    print(f"  Avg confidence  : {avg_conf:.2f}")
    print(f"  Time            : {elapsed}s")
    print(f"{'─'*60}")

    # Gate breakdown across all runs
    total_gates = {"auto-accept": 0, "review-required": 0, "quarantine": 0}
    for r in all_results:
        for gate, count in r.get("gate_counts", {}).items():
            if gate in total_gates:
                total_gates[gate] += count

    print(f"\n  Gate breakdown:")
    for gate, count in total_gates.items():
        pct = round(count / max(total_pages, 1) * 100, 1)
        flag = "  ← ACTION REQUIRED" if gate != "auto-accept" and count > 0 else ""
        print(f"    {gate:<20} {count:>4}  ({pct}%){flag}")

    print(f"\n  Next steps:")
    print(f"    1. Inspect ocr_raw/ outputs for quality")
    print(f"    2. Run 04_region_detector.py (Stage 5)")
    print(f"    3. Review quarantined pages manually\n")


if __name__ == "__main__":
    main()