src/lib/pdf_export/builder.py

"""PDF builders backed by PyMuPDF.

See the package docstring for the public API and tradeoffs.
"""

from __future__ import annotations

import json
import re
from collections.abc import Callable
from pathlib import Path

import fitz  # PyMuPDF

from src.config import load_config
from src.lib.storage import get_storage

# Bundled Arabic font for the searchable-PDF invisible-text layer.
# `insert_textbox(..., fontname='helv')` would silently encode every Arabic
# codepoint as '?' (Latin-1 only); without an Arabic-capable font the
# extracted text would be useless. NotoNaskhArabic is SIL-OFL licensed.
_FONTS_DIR = Path(__file__).resolve().parent / "fonts"
_AR_FONT_FILE = _FONTS_DIR / "NotoNaskhArabic-Regular.ttf"

# Strip OCR-stage scaffolding that the user almost certainly does NOT want
# read aloud: uncertainty markers and the [FOOTNOTE]…[/FOOTNOTE] wrappers.
# Keep the inner footnote text — that's still real prose. Drop only the
# delimiters.
_UNCERTAIN_RE = re.compile(r"\[\?([^\]]+)\]")
_FOOTNOTE_OPEN_RE = re.compile(r"\[FOOTNOTE\]")
_FOOTNOTE_CLOSE_RE = re.compile(r"\[/FOOTNOTE\]")


def _strip_markers(text: str) -> str:
    text = _UNCERTAIN_RE.sub(r"\1", text)         # [?word] → word
    text = _FOOTNOTE_OPEN_RE.sub("", text)
    text = _FOOTNOTE_CLOSE_RE.sub("", text)
    return text


def _page_text(book_id: str, page_num: int) -> str:
    """Best-available text for one page. Cleaned > OCR > empty.

    Returns a string with OCR markers stripped (they're noise for TTS).
    """
    cfg = load_config()
    clean_path = cfg.paths.clean_dir / book_id / f"page_{page_num:04d}.json"
    if clean_path.exists():
        data = json.loads(clean_path.read_text(encoding="utf-8"))
        return _strip_markers(data.get("text_clean") or data.get("text") or "")
    ocr_path = cfg.paths.ocr_dir / book_id / f"page_{page_num:04d}.json"
    if ocr_path.exists():
        data = json.loads(ocr_path.read_text(encoding="utf-8"))
        return _strip_markers(data.get("text") or "")
    return ""


def _html_escape(s: str) -> str:
    """Minimal HTML-entity escape for `insert_htmlbox`. We keep Arabic
    Unicode untouched; only escape the four characters HTML cares about.
    """
    return (
        s.replace("&", "&")
         .replace("<", "&lt;")
         .replace(">", "&gt;")
         .replace('"', "&quot;")
    )


def _text_to_html_block(text: str) -> str:
    paragraphs = [p.strip() for p in text.split("\n\n")]
    paragraphs = [p for p in paragraphs if p]
    if not paragraphs:
        return '<div dir="rtl" lang="ar"></div>'
    body = "".join(
        f"<p>{_html_escape(p).replace(chr(10), '<br>')}</p>"
        for p in paragraphs
    )
    return (
        '<div dir="rtl" lang="ar" '
        'style="font-size:13pt;line-height:1.85;text-align:right;">'
        f"{body}</div>"
    )


def make_text_pdf(book_id: str) -> bytes:
    """Build a fresh, text-only PDF for `book_id`. Returns the PDF bytes.

    One page per source page; each page renders the cleaned Arabic via
    `insert_htmlbox` (HarfBuzz, RTL-aware). Free, fast, small file —
    typically a few hundred KB for a 200-page book.

    Raises FileNotFoundError if neither OCR nor clean output exists for
    the book (caller should gate on `book.status >= 'ocr_done'`).
    """
    cfg = load_config()
    # Pull the per-page text bundle down from remote storage first (no-op on
    # the local backend) so the dir globs below see the cloud's OCR/clean JSON.
    get_storage().ensure_text_local(book_id)
    ocr_dir = cfg.paths.ocr_dir / book_id
    clean_dir = cfg.paths.clean_dir / book_id

    # Page set = union of OCR and clean dirs, sorted by page number.
    page_nums: set[int] = set()
    for d in (ocr_dir, clean_dir):
        if d.exists():
            for p in d.glob("page_*.json"):
                try:
                    page_nums.add(int(p.stem.split("_")[1]))
                except (IndexError, ValueError):
                    continue
    if not page_nums:
        raise FileNotFoundError(
            f"No OCR or cleaned page JSON for book_id={book_id!r}. "
            f"Run extraction first."
        )

    out = fitz.open()
    # A4 portrait — small file, predictable page size; the user is reading
    # by ear, not laying it out next to the original.
    width, height = 595.0, 842.0
    margin_x, margin_y = 40.0, 60.0
    rect = fitz.Rect(margin_x, margin_y, width - margin_x, height - margin_y)

    for n in sorted(page_nums):
        page = out.new_page(width=width, height=height)
        # Page-number footer in Western digits — a single small line so the
        # source page is still cite-able if the user prints this version.
        page.insert_textbox(
            fitz.Rect(margin_x, height - 50, width - margin_x, height - 30),
            f"page {n}",
            fontname="helv",
            fontsize=8,
            align=1,  # centered
        )
        text = _page_text(book_id, n) or "(no text on this page)"
        page.insert_htmlbox(rect, _text_to_html_block(text))

    # Without subsetting, insert_htmlbox embeds the full Arabic font on every
    # call — ~200 KB/page redundancy. subset_fonts() + garbage=4 deflate
    # gives a ~20× shrink on Arabic-heavy docs.
    out.subset_fonts()
    pdf_bytes: bytes = out.tobytes(garbage=4, deflate=True)
    out.close()
    return pdf_bytes


def make_searchable_pdf(book_id: str) -> bytes:
    """Original PDF + invisible per-page text layer. Returns PDF bytes.

    Page images are preserved exactly. We add the cleaned text via
    `insert_textbox(..., render_mode=3)` over each page — invisible but
    extractable by Cmd+F and screen readers.

    Caveat: Gemini OCR doesn't return word-level bounding boxes, so the
    invisible text is laid out as one block per page in document order
    rather than positioned per word. That's enough for TTS / search;
    it's NOT enough for "click word → highlight image". A bbox-aware OCR
    engine (Tesseract, Mistral OCR) is the upgrade path.

    Raises FileNotFoundError if the storage layer can't supply the
    original PDF, or if there's no extracted text yet.
    """
    if not _AR_FONT_FILE.exists():
        raise FileNotFoundError(
            f"Bundled Arabic font missing at {_AR_FONT_FILE}. "
            f"Run `python -m src.lib.pdf_export.fetch_font` to fetch it, "
            f"or copy NotoNaskhArabic-Regular.ttf into that directory."
        )

    store = get_storage()
    pdf_path: Path = store.ensure_local(book_id)
    store.ensure_text_local(book_id)  # materialise OCR/clean JSON (no-op locally)

    out = fitz.open(pdf_path)
    try:
        any_text = False
        for i in range(out.page_count):
            page_num = i + 1
            text = _page_text(book_id, page_num)
            if not text.strip():
                continue
            any_text = True
            page = out.load_page(i)
            # Insert invisible text covering the full page rect using a
            # bundled Arabic-capable TTF. Small font size so even a
            # render_mode-3 misbehavior couldn't cause visible glyph leak.
            # The text content stream carries the Unicode regardless of
            # the rendering mode (verified via round-trip: 73-char Arabic
            # passage extracts as the same Unicode it went in as).
            page.insert_textbox(
                page.rect,
                text,
                fontname="ar",
                fontfile=str(_AR_FONT_FILE),
                fontsize=6,
                render_mode=3,        # PDF spec: neither stroke nor fill
            )
        if not any_text:
            raise FileNotFoundError(
                f"No OCR/cleaned text on disk for book_id={book_id!r}. "
                f"Searchable PDF needs at least one extracted page."
            )
        # Same subsetting trick as the text-only builder — every page's
        # invisible-text insert can balloon the file otherwise.
        out.subset_fonts()
        return out.tobytes(garbage=4, deflate=True)
    finally:
        out.close()


def _pdf_export_cfg() -> dict:
    """Read the ``pdf_export`` config section with safe defaults.

    Kept here (not in :mod:`src.config`) because these knobs only matter
    to this module; a missing section yields the documented defaults so a
    fresh checkout works without editing YAML.
    """
    cfg = load_config().section("pdf_export")
    return {
        "strip_method": str(cfg.get("strip_method", "redact")).lower(),
        "rasterize_dpi": int(cfg.get("rasterize_dpi", 150)),
        "rasterize_format": str(cfg.get("rasterize_format", "jpeg")).lower(),
        "jpeg_quality": int(cfg.get("jpeg_quality", 80)),
        "only_rasterize_pages_with_text": bool(
            cfg.get("only_rasterize_pages_with_text", True)
        ),
    }


def make_clean_ocr_pdf(
    book_id: str,
    *,
    on_page: "Callable[[int, int], None] | None" = None,
) -> bytes:
    """Build a PDF where the OCR/clean text is the SOLE text layer.

    The problem this solves (and why it's distinct from
    :func:`make_searchable_pdf`): some source PDFs ship a *garbage*
    embedded text layer — a bad prior OCR pass, or mis-encoded glyphs.
    A read-aloud / TTS app reads that layer and produces gibberish.
    ``make_searchable_pdf`` only *adds* an invisible OCR layer on top and
    leaves the garbage in place, so the two text layers blend. Here we
    **destroy** the original text on any page that has one, then lay the
    clean OCR text under it invisibly. The result: the OCR is the single
    source of truth for every page's text, while the page still *looks*
    identical.

    How the original text is destroyed depends on ``pdf_export.strip_method``:

      * ``redact`` (default) — strip just the text layer via PyMuPDF
        redaction (``images=PDF_REDACT_IMAGE_NONE``), keeping the page's
        original embedded image untouched. No re-encode, so the file stays
        ~the same size as the source (often slightly smaller) and the build
        is fast. Ideal for scanned books, which is the whole corpus here.
      * ``rasterize`` — replace each text-layer page with a freshly rendered
        image at ``rasterize_dpi`` / ``rasterize_format``. Guarantees pixel
        fidelity for any page composition, but ~2× the file size and much
        slower. A fallback for pages where redaction wouldn't preserve the
        look (e.g. a born-digital page whose glyphs aren't backed by an
        image — rare for this corpus).

    Either way, pages with no embedded text layer keep their original
    content untouched and only gain the invisible OCR text.

    ``on_page(page_num, total)`` is invoked after each page so a job
    runner can emit progress. Returns the PDF bytes.

    Raises FileNotFoundError if the original PDF can't be materialised,
    the Arabic font is missing, or no page has any extracted text.
    """
    if not _AR_FONT_FILE.exists():
        raise FileNotFoundError(
            f"Bundled Arabic font missing at {_AR_FONT_FILE}. "
            f"Run `python -m src.lib.pdf_export.fetch_font` to fetch it."
        )

    opts = _pdf_export_cfg()
    store = get_storage()
    pdf_path: Path = store.ensure_local(book_id)
    store.ensure_text_local(book_id)  # materialise OCR/clean JSON (no-op locally)
    if opts["strip_method"] == "rasterize":
        return _build_rasterized(book_id, pdf_path, opts, on_page)
    return _build_redacted(book_id, pdf_path, on_page)


def _insert_invisible_text(page: "fitz.Page", text: str) -> None:
    """Lay `text` over the full page as render-mode-3 (invisible) Arabic."""
    page.insert_textbox(
        page.rect,
        text,
        fontname="ar",
        fontfile=str(_AR_FONT_FILE),
        fontsize=6,
        render_mode=3,            # PDF spec: neither stroke nor fill
    )


def _build_redacted(
    book_id: str,
    pdf_path: Path,
    on_page: "Callable[[int, int], None] | None",
) -> bytes:
    """strip_method='redact': remove the text layer in place, keep the image.

    Works on the original document: for each page that has a text layer we
    redact it away (images preserved), then overlay the invisible OCR text.
    Pages with no text layer are left as-is plus the OCR overlay.
    """
    doc = fitz.open(pdf_path)
    try:
        total = doc.page_count
        any_text = False
        for i in range(total):
            page = doc.load_page(i)
            if page.get_text("text").strip():
                # Full-page redaction strips every text span; PDF_REDACT_IMAGE_NONE
                # leaves embedded images (the actual scan) untouched, so the page
                # still looks identical and the file doesn't grow.
                page.add_redact_annot(page.rect)
                page.apply_redactions(images=fitz.PDF_REDACT_IMAGE_NONE)
            text = _page_text(book_id, i + 1)
            if text.strip():
                any_text = True
                _insert_invisible_text(page, text)
            if on_page is not None:
                on_page(i + 1, total)
        if not any_text:
            raise FileNotFoundError(
                f"No OCR/cleaned text on disk for book_id={book_id!r}. "
                f"Run extraction before building the clean-OCR PDF."
            )
        doc.subset_fonts()
        return doc.tobytes(garbage=4, deflate=True)
    finally:
        doc.close()


def _build_rasterized(
    book_id: str,
    pdf_path: Path,
    opts: dict,
    on_page: "Callable[[int, int], None] | None",
) -> bytes:
    """strip_method='rasterize': flatten text-layer pages to an image.

    Builds a fresh document. Pages with a text layer (or all pages when
    ``only_rasterize_pages_with_text`` is false) are rendered to an image —
    which carries no text stream, so the original text is gone. Pure-scan
    pages are copied through untouched to avoid a needless re-encode.
    """
    dpi = opts["rasterize_dpi"]
    fmt = opts["rasterize_format"]
    quality = opts["jpeg_quality"]
    only_with_text = opts["only_rasterize_pages_with_text"]

    src = fitz.open(pdf_path)
    out = fitz.open()
    try:
        total = src.page_count
        any_text = False
        for i in range(total):
            src_page = src.load_page(i)
            has_text = bool(src_page.get_text("text").strip())
            if has_text or not only_with_text:
                pix = src_page.get_pixmap(dpi=dpi, alpha=False)
                img_bytes = (
                    pix.tobytes(output="png")
                    if fmt == "png"
                    else pix.tobytes(output="jpeg", jpg_quality=quality)
                )
                new_page = out.new_page(
                    width=src_page.rect.width, height=src_page.rect.height
                )
                new_page.insert_image(new_page.rect, stream=img_bytes)
            else:
                out.insert_pdf(src, from_page=i, to_page=i)
                new_page = out.load_page(out.page_count - 1)

            text = _page_text(book_id, i + 1)
            if text.strip():
                any_text = True
                _insert_invisible_text(new_page, text)
            if on_page is not None:
                on_page(i + 1, total)

        if not any_text:
            raise FileNotFoundError(
                f"No OCR/cleaned text on disk for book_id={book_id!r}. "
                f"Run extraction before building the clean-OCR PDF."
            )
        out.subset_fonts()
        return out.tobytes(garbage=4, deflate=True)
    finally:
        out.close()
        src.close()