src/lib/metadata_probe.py

"""Smart metadata probe for the Add Book wizard.

Given the bytes of an uploaded PDF, this module:
  1. Counts pages (free, PyMuPDF).
  2. Detects native text vs scan via a chars-per-page peek (free, PyMuPDF).
  3. Sends the first N pages to Gemini for a structured-JSON metadata guess —
     as TEXT for native-text PDFs (cheaper, more accurate against the real
     text layer) or as page IMAGES for scans.

Output is a typed dataclass shaped to prefill the form. The probe runs only
when the user clicks "Analyze" in the UI; results are cached in session_state
keyed by SHA-256 of the upload bytes so form re-renders never re-bill Gemini.
"""

from __future__ import annotations

import json
import re
from dataclasses import asdict, dataclass, field
from pathlib import Path

import fitz  # PyMuPDF

from src.config import REPO_ROOT, load_config
from src.stage3_ocr.gemini_client import generate, image_part

PROMPT_PATH = REPO_ROOT / "prompts" / "metadata_probe.txt"

ERAS = ("ante-nicene", "nicene", "post-nicene", "medieval", "modern")
TRADITIONS = ("coptic", "antiochene", "cappadocian", "latin", "syriac", "other")
LANGUAGES = ("ar", "en", "gr", "la")
RELIGIONS = ("christian", "islamic", "jewish", "other")

# Average chars/page on the first 3 pages above which we treat the PDF as a
# real text-layer PDF (vs a scanned image PDF). 100 is well below the density
# of any typeset Arabic page and well above what stray OCR junk produces.
NATIVE_TEXT_MIN_CHARS_PER_PAGE = 100


@dataclass
class ProbeResult:
    """One analyze() call's output, ready to prefill the Add Book form."""
    title_ar: str | None
    title_en: str | None
    author: str | None
    author_id: str | None
    era: str | None
    tradition: str | None
    language: str | None
    pages_total: int
    extraction_mode_suggested: str  # "ocr" | "native_text"
    avg_chars_per_page: int         # what the native-text peek measured
    confidence: str                 # "low" | "medium" | "high"
    rationale: str | None
    suggested_labels: list[str] = field(default_factory=list)
    book_religion: str | None = None
    author_religion: str | None = None
    error: str | None = None        # set if Gemini call or JSON parse failed

    def asdict(self) -> dict:
        return asdict(self)


# ---------- Free probes (no API) -------------------------------------------

def page_count(pdf_bytes: bytes) -> int:
    try:
        with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
            return doc.page_count
    except Exception:  # noqa: BLE001
        return 0


def peek_native_text(pdf_bytes: bytes, n_pages: int = 3) -> tuple[bool, int]:
    """(is_text_pdf, avg_chars/page on the first n pages).

    Returns (False, 0) on parse failure so callers default to OCR — the
    safer choice for an unreadable PDF.
    """
    try:
        with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
            n = min(n_pages, doc.page_count)
            if n == 0:
                return False, 0
            total = sum(
                len(doc.load_page(i).get_text("text").strip())
                for i in range(n)
            )
            avg = total // n
            return avg >= NATIVE_TEXT_MIN_CHARS_PER_PAGE, avg
    except Exception:  # noqa: BLE001
        return False, 0


# ---------- Sample helpers --------------------------------------------------

def sample_text(pdf_bytes: bytes, n_pages: int) -> str:
    chunks: list[str] = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        n = min(n_pages, doc.page_count)
        for i in range(n):
            text = doc.load_page(i).get_text("text").strip()
            chunks.append(f"=== Page {i+1} ===\n{text}")
    return "\n\n".join(chunks)


def _sample_images(pdf_bytes: bytes, n_pages: int, dpi: int = 150) -> list[bytes]:
    zoom = dpi / 72
    matrix = fitz.Matrix(zoom, zoom)
    out: list[bytes] = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        n = min(n_pages, doc.page_count)
        for i in range(n):
            pix = doc.load_page(i).get_pixmap(matrix=matrix, alpha=False)
            out.append(pix.tobytes("png"))
    return out


# ---------- JSON parsing helpers --------------------------------------------

def _existing_label_names() -> list[str]:
    """Manual labels already in the user's library, for the prompt context.

    Lazy import to avoid a labels → metadata_probe coupling at module load.
    Returns [] on any failure — the prompt template tolerates an empty list.
    """
    try:
        from src.lib.labels import list_labels
        return [
            (l.display_en or l.display_ar)
            for l in list_labels(kind="manual")
            if (l.display_en or l.display_ar)
        ]
    except Exception:  # noqa: BLE001
        return []


def _fill_prompt_template(prompt: str) -> str:
    """Fill `{existing_labels_list}` with the user's known manual labels."""
    names = _existing_label_names()
    listed = ", ".join(sorted(names)) if names else "(none yet — propose new ones if any apply)"
    return prompt.replace("{existing_labels_list}", listed)


def _strip_fences(s: str) -> str:
    s = s.strip()
    if s.startswith("```"):
        s = re.sub(r"^```[a-zA-Z]*\n?", "", s)
        s = re.sub(r"\n?```\s*$", "", s)
    return s.strip()


def _pick_enum(value, allowed: tuple[str, ...]) -> str | None:
    if not value:
        return None
    v = str(value).strip().lower()
    return v if v in allowed else None


def _clean_suggested_labels(value, *, max_labels: int = 5) -> list[str]:
    """Coerce Gemini's suggested_labels into a clean list of short strings.

    Filters out non-strings, lowercases, dedupes (preserving order), trims
    to <=64 chars per tag, caps the list length, and drops anything that's
    obviously era/tradition/language (those have their own fields, even
    though the prompt asks the model not to duplicate them).
    """
    if not isinstance(value, list):
        return []
    excluded = set(ERAS) | set(TRADITIONS) | set(LANGUAGES)
    out: list[str] = []
    seen: set[str] = set()
    for s in value:
        if not isinstance(s, str):
            continue
        v = s.strip().lower()[:64]
        if not v or v in seen or v in excluded:
            continue
        seen.add(v)
        out.append(v)
        if len(out) >= max_labels:
            break
    return out


def _normalize(
    raw: dict,
    *,
    pages_total: int,
    extraction_mode_suggested: str,
    avg_chars_per_page: int,
) -> ProbeResult:
    return ProbeResult(
        title_ar=(raw.get("title_ar") or None),
        title_en=(raw.get("title_en") or None),
        author=(raw.get("author") or None),
        author_id=(raw.get("author_id") or None),
        era=_pick_enum(raw.get("era"), ERAS),
        tradition=_pick_enum(raw.get("tradition"), TRADITIONS),
        language=_pick_enum(raw.get("language"), LANGUAGES),
        pages_total=pages_total,
        extraction_mode_suggested=extraction_mode_suggested,
        avg_chars_per_page=avg_chars_per_page,
        confidence=str(raw.get("confidence") or "low").lower(),
        rationale=(raw.get("rationale") or None),
        suggested_labels=_clean_suggested_labels(raw.get("suggested_labels")),
        book_religion=_pick_enum(raw.get("book_religion"), RELIGIONS),
        author_religion=_pick_enum(raw.get("author_religion"), RELIGIONS),
    )


# ---------- Public entry point ---------------------------------------------

def _ask_gemini(parts: list, *, model: str, note: str) -> tuple[dict, str | None]:
    """One Gemini structured-JSON call. Returns (parsed_dict, error_or_none)."""
    raw_text = ""
    try:
        raw_text = generate(
            model=model,
            parts=parts,
            response_mime_type="application/json",
            stage="metadata_probe",
            book_id=None,
            note=note,
        )
    except Exception as e:  # noqa: BLE001
        return {}, f"{type(e).__name__}: {e}"
    if not raw_text:
        return {}, None
    try:
        parsed = json.loads(_strip_fences(raw_text))
    except json.JSONDecodeError as e:
        return {}, f"JSON parse: {e}"
    return (parsed, None) if isinstance(parsed, dict) else ({}, "model returned non-object JSON")


def probe_metadata_from_text(
    text: str,
    *,
    pages_total: int,
    extraction_mode_suggested: str = "ocr",
    avg_chars_per_page: int = 0,
) -> ProbeResult:
    """Ask Gemini for metadata given a chunk of already-extracted text.

    Used by the OCR sample flow: after we OCR + clean N pages of an unsaved
    PDF, the cleaned text is fed back through the same metadata prompt as
    the text-layer path. Same prompt, same enum normalization, same shape
    of `ProbeResult` — keeping a single source of truth for what "metadata"
    means downstream.
    """
    cfg = load_config()
    section = cfg.section("metadata_probe")
    model = section.get("model", "gemini-2.5-flash")
    prompt = _fill_prompt_template(PROMPT_PATH.read_text(encoding="utf-8"))

    raw_dict, error = _ask_gemini(
        [
            prompt,
            "\n\n--- Sample text from the first pages "
            "(produced by OCR + cleanup of the uploaded PDF) ---\n\n"
            + (text or ""),
        ],
        model=model,
        note=f"source=ocr_sample chars={len(text or '')}",
    )
    result = _normalize(
        raw_dict,
        pages_total=pages_total,
        extraction_mode_suggested=extraction_mode_suggested,
        avg_chars_per_page=avg_chars_per_page,
    )
    result.error = error
    return result


def probe_metadata(
    pdf_bytes: bytes,
    *,
    force_mode: str | None = None,
) -> ProbeResult:
    """Detect mode, sample pages, ask Gemini, return a normalized result.

    By default the probe uses the auto-detect peek to choose between text
    sampling (native_text PDFs) and image sampling (scans). Pass
    `force_mode="ocr"` to send page images even when the PDF has a text
    layer — useful when the embedded text is corrupted or the title
    metadata lives only on the cover image.

    `force_mode` only affects this one metadata call. The
    `extraction_mode_suggested` field of the result still reflects the
    PDF's own text-layer peek (i.e. what the actual ingest should do).
    """
    cfg = load_config()
    section = cfg.section("metadata_probe")
    model = section.get("model", "gemini-2.5-flash")
    n_pages = int(section.get("pages_to_sample", 5))

    pages = page_count(pdf_bytes)
    is_native_pdf, avg_chars = peek_native_text(pdf_bytes)
    suggested_mode = "native_text" if is_native_pdf else "ocr"

    # Pick how to sample for THIS probe call. Forcing native_text on a PDF
    # without a text layer would send empty content to Gemini, so fall back
    # to images in that case.
    if force_mode == "ocr":
        probe_mode = "ocr"
    elif force_mode == "native_text" and is_native_pdf:
        probe_mode = "native_text"
    else:
        probe_mode = "native_text" if is_native_pdf else "ocr"

    prompt = _fill_prompt_template(PROMPT_PATH.read_text(encoding="utf-8"))

    parts: list = [prompt]
    if probe_mode == "native_text":
        parts.append(
            "\n\n--- Sample text from the first pages "
            "(extracted via PyMuPDF, exactly as embedded in the PDF) ---\n\n"
            + sample_text(pdf_bytes, n_pages)
        )
    else:
        for png in _sample_images(pdf_bytes, n_pages):
            parts.append(image_part(png))

    raw_dict, error = _ask_gemini(
        parts,
        model=model,
        note=(
            f"probe_mode={probe_mode} "
            f"forced={bool(force_mode)} "
            f"pages_sampled={min(n_pages, pages)}"
        ),
    )
    result = _normalize(
        raw_dict,
        pages_total=pages,
        extraction_mode_suggested=suggested_mode,
        avg_chars_per_page=avg_chars,
    )
    result.error = error
    return result