src/api/routers/samples.py

"""OCR-sample preview for the Add Book wizard.

Before committing to a full ingest, the wizard can OCR + clean the first N
pages of an uploaded PDF and show a per-page *image | native text | OCR |
cleaned* comparison, so the user decides the extraction mode with eyes open.

- ``POST /books/ocr-sample`` body ``{uploadId, pages}`` → 202 ``{jobId, sseUrl}``.
  Runs the ``ocr_sample`` job (paid; SSE-tracked). Admin-only.
- ``GET /uploads/{uploadId}/pages/{n}/image?dpi=`` → PNG of one page of the
  *uploaded* (not-yet-ingested) PDF, rendered on demand.
- ``GET /uploads/{uploadId}/ocr-sample`` → per-page OCR + cleaned text read from
  the job's scratch ``__probe__{uploadId[:12]}`` dirs. Empty ``pages`` until the
  job has run.

The scratch dirs are keyed by ``upload_id`` (see ``jobs/types/ocr_sample.py``),
so these reads need only the upload id — no sha/probe-id indirection on the FE.
"""
from __future__ import annotations

import json
import logging
import uuid
from pathlib import Path as PathLib

from fastapi import APIRouter, Depends, Path, Query
from fastapi.responses import Response

from src.api.deps import require_admin
from src.api.dto.common import ApiModel, JobStartResponse
from src.api.errors import BadRequest
from src.api.jobs import runner, store
from src.api.limits import enforce_paid_request
from src.api.validation import require_safe_id
from src.config import load_config
from src.lib.auth.models import User
from src.lib.probe_ocr import DEFAULT_SAMPLE_PAGES, MAX_SAMPLE_PAGES, probe_id_for

# Side-effect import: registers the "ocr_sample" handler with the job registry.
from src.api.jobs.types import ocr_sample as _ocr_sample_job  # noqa: F401

logger = logging.getLogger(__name__)

router = APIRouter(tags=["samples"])


class OcrSampleRequest(ApiModel):
    """Body for ``POST /books/ocr-sample``."""

    upload_id: str
    pages: int = DEFAULT_SAMPLE_PAGES


def _upload_path(upload_id: str) -> PathLib:
    # Validate before the join so a crafted id can't escape the uploads dir.
    require_safe_id(upload_id, field="uploadId")
    return load_config().paths.data_dir / "uploads" / f"{upload_id}.pdf"


@router.post("/books/ocr-sample", status_code=202, response_model=JobStartResponse)
def post_ocr_sample(
    payload: OcrSampleRequest,
    user: User = Depends(require_admin),
) -> JobStartResponse:
    """Kick off an OCR-sample job for an uploaded PDF. Returns 202 + job."""
    enforce_paid_request(user)  # daily spend ceiling + per-user rate limit
    if not _upload_path(payload.upload_id).exists():
        raise BadRequest(
            f"No upload with id {payload.upload_id!r}. Re-upload and try again."
        )
    pages = max(1, min(int(payload.pages or DEFAULT_SAMPLE_PAGES), MAX_SAMPLE_PAGES))
    params = {"upload_id": payload.upload_id, "pages": pages}
    job_id = uuid.uuid4().hex
    store.insert_job(
        id=job_id,
        type="ocr_sample",
        params_json=params,
        username=user.username,
        subject_id=payload.upload_id,
    )
    runner.enqueue(
        job_id,
        type_="ocr_sample",
        params=params,
        subject_key=f"probe:{payload.upload_id}",
    )
    logger.info("ocr_sample enqueued: job=%s upload=%s pages=%s", job_id, payload.upload_id, pages)
    return JobStartResponse(job_id=job_id, sse_url=f"/jobs/{job_id}/events")


@router.get("/uploads/{uploadId}/pages/{n}/image")
def get_upload_page_image(
    upload_id: str = Path(..., alias="uploadId"),
    n: int = Path(..., ge=1),
    dpi: int = Query(default=150, ge=72, le=300),
    _: User = Depends(require_admin),
) -> Response:
    """PNG of one page of an uploaded (pre-ingest) PDF. Rendered on demand."""
    path = _upload_path(upload_id)
    if not path.exists():
        raise BadRequest(f"No upload with id {upload_id!r}.")
    import fitz  # local import keeps PyMuPDF out of module import cost

    try:
        with fitz.open(path) as doc:
            if n > doc.page_count:
                raise BadRequest(f"Upload {upload_id!r} has only {doc.page_count} pages.")
            zoom = dpi / 72
            pix = doc.load_page(n - 1).get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
            png = pix.tobytes("png")
    except BadRequest:
        raise
    except Exception as e:  # noqa: BLE001
        raise BadRequest(f"Could not render page {n} of upload {upload_id!r}: {e}")
    return Response(content=png, media_type="image/png",
                    headers={"Cache-Control": "private, max-age=3600"})


@router.get("/uploads/{uploadId}/ocr-sample")
def get_ocr_sample(
    upload_id: str = Path(..., alias="uploadId"),
    _: User = Depends(require_admin),
) -> dict:
    """Per-page OCR + cleaned text from the sample job's scratch dirs.

    Returns ``{uploadId, pages: [...]}``. ``pages`` is empty until the
    ``ocr_sample`` job has produced scratch output. camelCase keys are
    written directly (no DTO round-trip needed).
    """
    require_safe_id(upload_id, field="uploadId")
    cfg = load_config()
    probe_id = probe_id_for(upload_id)
    ocr_dir = cfg.paths.ocr_dir / probe_id
    clean_dir = cfg.paths.clean_dir / probe_id

    page_nums: set[int] = set()
    for d in (ocr_dir, clean_dir):
        if d.exists():
            for f in d.glob("page_*.json"):
                try:
                    page_nums.add(int(f.stem.split("_")[1]))
                except (IndexError, ValueError):
                    continue

    pages: list[dict] = []
    for n in sorted(page_nums):
        ocr_blob: dict = {}
        clean_blob: dict = {}
        op = ocr_dir / f"page_{n:04d}.json"
        cp = clean_dir / f"page_{n:04d}.json"
        if op.exists():
            ocr_blob = json.loads(op.read_text(encoding="utf-8"))
        if cp.exists():
            clean_blob = json.loads(cp.read_text(encoding="utf-8"))
        pages.append({
            "pdfPage": n,
            "ocrText": ocr_blob.get("text") or "",
            "cleanText": clean_blob.get("text_clean") or "",
            "uncertainWords": list(ocr_blob.get("uncertain_words") or []),
            "cleanupWarnings": list(clean_blob.get("cleanup_warnings") or []),
            "ocrEngine": ocr_blob.get("ocr_engine"),
            "cleanupEngine": clean_blob.get("cleanup_engine"),
        })

    return {"uploadId": upload_id, "pages": pages}