src/lib/probe_ocr.py

"""OCR + cleanup sample run for the Add Book wizard.

Runs the standard Stage 3 OCR + Stage 4 cleanup logic on a chosen subset of
pages from an unsaved PDF, without first inserting the book into inventory.

Output lives under `data/{ocr,clean}/__probe__{sha[:12]}/` so that:
- The user can preview OCR + cleanup quality on a real sample before
  committing to a full ingest. If quality is bad they can abort with no
  cost beyond the sample.
- On Submit, those scratch directories are merged into the actual
  `data/{ocr,clean}/{book_id}/` so the standard ingest pipeline (which is
  per-page idempotent) skips them on its OCR/cleanup passes — sample spend
  rolls forward into the real ingest.

Cost is recorded under the synthetic probe_id so per-book accounting in
Costs & usage stays accurate even when the book is later renamed.
"""

from __future__ import annotations

import json
import shutil
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

import fitz  # PyMuPDF

from src.config import load_config
from src.lib.costs import BudgetExceededError
from src.lib.job_context import JobCanceledError
from src.stage3_ocr.gemini_client import RateLimiter, generate, image_part
from src.stage3_ocr.parsers import extract_footnotes, extract_uncertain_words
from src.stage3_ocr.run_ocr import OCR_PROMPT
from src.stage4_cleanup.run_cleanup import CLEANUP_PROMPT

PROBE_PREFIX = "__probe__"
MAX_SAMPLE_PAGES = 50
DEFAULT_SAMPLE_PAGES = 10


@dataclass
class PageSample:
    """One probed page: original image + OCR output + cleanup output."""
    page_num: int
    original_png: bytes
    ocr_text: str = ""
    ocr_uncertain: list[str] = field(default_factory=list)
    ocr_footnotes: list[str] = field(default_factory=list)
    cleaned_text: str = ""
    cleanup_warnings: list[str] = field(default_factory=list)
    error: str | None = None


@dataclass
class OcrSampleResult:
    """One run_ocr_sample() output, ready for the comparison view."""
    probe_id: str            # synthetic __probe__{sha[:12]}
    pages_total: int         # pages in the source PDF
    pages: list[PageSample]
    metadata_text: str       # concatenated cleaned text for metadata extraction


def probe_id_for(sha: str) -> str:
    return f"{PROBE_PREFIX}{sha[:12]}"


# ---------- Free helpers (no API) ------------------------------------------

def _render_page_png(pdf_bytes: bytes, page_num: int, dpi: int) -> bytes:
    zoom = dpi / 72
    matrix = fitz.Matrix(zoom, zoom)
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        page = doc.load_page(page_num - 1)
        pix = page.get_pixmap(matrix=matrix, alpha=False)
        return pix.tobytes("png")


def _shrink_warnings(orig: str, cleaned: str) -> list[str]:
    """Same fidelity guard the main cleanup driver uses."""
    warnings: list[str] = []
    if orig:
        shrink = 1 - (len(cleaned) / len(orig))
        if shrink > 0.10:
            warnings.append(f"shrunk {shrink*100:.0f}%")
    orig_fn  = orig.count("[FOOTNOTE]")
    clean_fn = cleaned.count("[FOOTNOTE]")
    if clean_fn < orig_fn:
        warnings.append(f"lost {orig_fn - clean_fn} [FOOTNOTE] block(s)")
    return warnings


# ---------- API calls -------------------------------------------------------

def _ocr_one(model: str, png: bytes, rl: RateLimiter, *, probe_id: str, page_num: int) -> str:
    # book_id=None to avoid the FOREIGN KEY constraint on llm_calls.book_id
    # (the books table doesn't have a row for the synthetic __probe__ id and
    # we'd rather not insert+rollback per probe). The probe_id lives in `note`
    # so the calls are still traceable in Operations → Costs & usage.
    return generate(
        model=model,
        parts=[image_part(png), OCR_PROMPT],
        rate_limiter=rl,
        stage="ocr",
        book_id=None,
        note=f"probe={probe_id} page={page_num}",
    )


def _cleanup_one(model: str, ocr_text: str, rl: RateLimiter, *, probe_id: str, page_num: int) -> str:
    return generate(
        model=model,
        parts=[CLEANUP_PROMPT + ocr_text],
        rate_limiter=rl,
        stage="cleanup",
        book_id=None,
        note=f"probe={probe_id} page={page_num}",
    )


# ---------- Public entry points --------------------------------------------

def run_ocr_sample(
    pdf_bytes: bytes,
    *,
    sha: str,
    page_numbers: list[int],
) -> OcrSampleResult:
    """OCR + cleanup the chosen pages, persist per-page JSON to scratch dirs.

    Continues on per-page failure — the user gets to see which pages OCR'd
    cleanly and which didn't. Render failures, OCR failures, and cleanup
    failures each set `PageSample.error` and skip the downstream step.
    """
    cfg = load_config()
    ocr_cfg     = cfg.section("ocr")
    cleanup_cfg = cfg.section("cleanup")
    ocr_model     = ocr_cfg["primary"]
    cleanup_model = cleanup_cfg["model"]
    pdf_dpi       = int(ocr_cfg.get("pdf_dpi", 300))
    rate_limit    = int(ocr_cfg.get("rate_limit_per_minute", 60))

    if len(page_numbers) > MAX_SAMPLE_PAGES:
        raise ValueError(
            f"sample size {len(page_numbers)} exceeds cap of {MAX_SAMPLE_PAGES}"
        )

    probe_id  = probe_id_for(sha)
    ocr_dir   = cfg.paths.ocr_dir   / probe_id
    clean_dir = cfg.paths.clean_dir / probe_id
    ocr_dir.mkdir(parents=True, exist_ok=True)
    clean_dir.mkdir(parents=True, exist_ok=True)

    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        total = doc.page_count
    valid = sorted({p for p in page_numbers if 1 <= p <= total})
    if not valid:
        return OcrSampleResult(
            probe_id=probe_id, pages_total=total, pages=[], metadata_text="",
        )

    rl_ocr     = RateLimiter(max_per_minute=rate_limit)
    rl_cleanup = RateLimiter(max_per_minute=rate_limit)

    samples: list[PageSample] = []
    cleaned_chunks: list[str] = []

    for n in valid:
        # Always render the original — used in the preview pane regardless of
        # whether OCR succeeds.
        try:
            png = _render_page_png(pdf_bytes, n, pdf_dpi)
        except Exception as e:  # noqa: BLE001
            samples.append(PageSample(
                page_num=n, original_png=b"",
                error=f"render failed: {type(e).__name__}: {e}",
            ))
            continue

        # OCR
        try:
            ocr_text = _ocr_one(ocr_model, png, rl_ocr, probe_id=probe_id, page_num=n)
        except (JobCanceledError, BudgetExceededError):
            # A cancel or tripped spend cap must abort the whole sample job, not
            # be recorded as one page's error while the loop keeps going.
            raise
        except Exception as e:  # noqa: BLE001
            samples.append(PageSample(
                page_num=n, original_png=png,
                error=f"OCR failed: {type(e).__name__}: {e}",
            ))
            continue

        ocr_payload = {
            "book_id":         probe_id,
            "page_number":     n,
            "language":        "ar",
            "text":            ocr_text,
            "uncertain_words": extract_uncertain_words(ocr_text),
            "footnotes":       extract_footnotes(ocr_text),
            "ocr_engine":      ocr_model,
            "ocr_timestamp":   datetime.now(timezone.utc).isoformat(),
            "raw_response":    ocr_text,
        }
        (ocr_dir / f"page_{n:04d}.json").write_text(
            json.dumps(ocr_payload, ensure_ascii=False, indent=2), encoding="utf-8",
        )

        # Cleanup
        try:
            clean_text = _cleanup_one(cleanup_model, ocr_text, rl_cleanup, probe_id=probe_id, page_num=n)
        except (JobCanceledError, BudgetExceededError):
            raise  # cancel/budget aborts the whole sample job
        except Exception as e:  # noqa: BLE001
            samples.append(PageSample(
                page_num=n,
                original_png=png,
                ocr_text=ocr_text,
                ocr_uncertain=ocr_payload["uncertain_words"],
                ocr_footnotes=ocr_payload["footnotes"],
                error=f"cleanup failed: {type(e).__name__}: {e}",
            ))
            continue

        warnings = _shrink_warnings(ocr_text, clean_text)
        clean_payload = dict(ocr_payload)
        clean_payload.update({
            "text_clean":            clean_text,
            "cleanup_engine":        cleanup_model,
            "cleanup_timestamp":     datetime.now(timezone.utc).isoformat(),
            "uncertain_words_clean": extract_uncertain_words(clean_text),
            "cleanup_warnings":      warnings,
        })
        (clean_dir / f"page_{n:04d}.json").write_text(
            json.dumps(clean_payload, ensure_ascii=False, indent=2), encoding="utf-8",
        )

        samples.append(PageSample(
            page_num=n,
            original_png=png,
            ocr_text=ocr_text,
            ocr_uncertain=ocr_payload["uncertain_words"],
            ocr_footnotes=ocr_payload["footnotes"],
            cleaned_text=clean_text,
            cleanup_warnings=warnings,
        ))
        cleaned_chunks.append(f"=== Page {n} ===\n{clean_text}")

    return OcrSampleResult(
        probe_id=probe_id,
        pages_total=total,
        pages=samples,
        metadata_text="\n\n".join(cleaned_chunks),
    )


def carry_forward_to(book_id: str, *, sha: str) -> tuple[int, int]:
    """Move scratch __probe__ OCR/clean dirs into the actual book_id dirs.

    Returns (pages_carried_ocr, pages_carried_clean). Existing destination
    files are left untouched — same per-page idempotency the main pipeline
    relies on.
    """
    cfg = load_config()
    probe_id = probe_id_for(sha)
    moved_ocr   = _merge_dir(cfg.paths.ocr_dir   / probe_id, cfg.paths.ocr_dir   / book_id)
    moved_clean = _merge_dir(cfg.paths.clean_dir / probe_id, cfg.paths.clean_dir / book_id)
    for d in (cfg.paths.ocr_dir / probe_id, cfg.paths.clean_dir / probe_id):
        if d.exists() and not any(d.iterdir()):
            try:
                d.rmdir()
            except OSError:
                pass
    return moved_ocr, moved_clean


def _merge_dir(src: Path, dst: Path) -> int:
    if not src.exists():
        return 0
    dst.mkdir(parents=True, exist_ok=True)
    moved = 0
    for f in src.iterdir():
        target = dst / f.name
        if target.exists():
            continue
        shutil.move(str(f), str(target))
        moved += 1
    return moved


def discard_probe(sha: str) -> None:
    """Delete the scratch __probe__ dirs. Use when the user picks a
    different file or aborts the wizard without saving."""
    cfg = load_config()
    probe_id = probe_id_for(sha)
    for d in (cfg.paths.ocr_dir / probe_id, cfg.paths.clean_dir / probe_id):
        if d.exists():
            shutil.rmtree(d, ignore_errors=True)