document_processing.py

from __future__ import annotations

import io
import os
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import pypdfium2 as pdfium
from PIL import Image, UnidentifiedImageError

MAX_DOCUMENT_BYTES = 20 * 1024 * 1024
MAX_PAGES = 10
MAX_IMAGE_EDGE = 2048
RENDER_DPI = 200
PNG_COMPRESS_LEVEL = 6
MAX_IMAGE_PIXELS = 50_000_000

Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS


class DocumentError(ValueError):
    """A safe, user-facing document validation error."""


@dataclass(frozen=True)
class PreparedPage:
    number: int
    png: bytes
    width: int
    height: int


@dataclass(frozen=True)
class PreparedDocument:
    name: str
    size: int
    kind: str
    total_pages: int
    pages: tuple[PreparedPage, ...]

    @property
    def processed_pages(self) -> int:
        return len(self.pages)

    def to_session(self) -> dict:
        return {
            "name": self.name,
            "size": self.size,
            "kind": self.kind,
            "total_pages": self.total_pages,
            "pages": [
                {
                    "number": page.number,
                    "png": page.png,
                    "width": page.width,
                    "height": page.height,
                }
                for page in self.pages
            ],
            "current_page": 0,
            "results": [None] * len(self.pages),
        }


def _read_prefix(path: Path, length: int = 1024) -> bytes:
    with path.open("rb") as handle:
        return handle.read(length)


def detect_document_kind(path: Path) -> str:
    prefix = _read_prefix(path)
    if prefix.startswith(b"%PDF-"):
        return "pdf"

    try:
        with Image.open(path) as image:
            image.verify()
            image_format = (image.format or "").upper()
    except (UnidentifiedImageError, OSError, SyntaxError) as exc:
        raise DocumentError("Use a valid PDF, PNG, JPEG, or WebP file.") from exc

    supported = {"PNG": "png", "JPEG": "jpeg", "WEBP": "webp"}
    if image_format not in supported:
        raise DocumentError("Use a PDF, PNG, JPEG, or WebP file.")
    return supported[image_format]


def _validate_upload_location(path: Path) -> None:
    roots = {Path(tempfile.gettempdir()).resolve()}
    configured = os.environ.get("GRADIO_TEMP_DIR")
    if configured:
        roots.add(Path(configured).expanduser().resolve())

    resolved = path.resolve()
    if not any(resolved == root or root in resolved.parents for root in roots):
        raise DocumentError("The upload could not be validated.")


def _encode_png(image: Image.Image) -> PreparedPage:
    image = image.convert("RGB")
    image.thumbnail((MAX_IMAGE_EDGE, MAX_IMAGE_EDGE), Image.Resampling.LANCZOS)
    buffer = io.BytesIO()
    image.save(buffer, format="PNG", compress_level=PNG_COMPRESS_LEVEL)
    return PreparedPage(
        number=0,
        png=buffer.getvalue(),
        width=image.width,
        height=image.height,
    )


def _prepare_image(path: Path) -> tuple[PreparedPage, ...]:
    try:
        with Image.open(path) as image:
            if getattr(image, "is_animated", False):
                image.seek(0)
            prepared = _encode_png(image)
    except Image.DecompressionBombError as exc:
        raise DocumentError("This image is too large to process safely.") from exc
    except (UnidentifiedImageError, OSError, SyntaxError) as exc:
        raise DocumentError("The image could not be decoded.") from exc

    return (
        PreparedPage(
            number=1,
            png=prepared.png,
            width=prepared.width,
            height=prepared.height,
        ),
    )


def _render_pdf_page(document: pdfium.PdfDocument, index: int) -> PreparedPage:
    page = document[index]
    bitmap = None
    try:
        width, height = page.get_size()
        if width <= 0 or height <= 0:
            raise DocumentError(f"Page {index + 1} has invalid dimensions.")

        scale = min(
            RENDER_DPI / 72,
            MAX_IMAGE_EDGE / max(width, height),
        )
        bitmap = page.render(scale=max(scale, 0.1))
        image = bitmap.to_pil()
        prepared = _encode_png(image)
        return PreparedPage(
            number=index + 1,
            png=prepared.png,
            width=prepared.width,
            height=prepared.height,
        )
    except DocumentError:
        raise
    except Exception as exc:
        raise DocumentError(f"Page {index + 1} could not be rendered.") from exc
    finally:
        if bitmap is not None:
            bitmap.close()
        page.close()


def _prepare_pdf(
    path: Path,
    on_progress: Callable[[int, int], None] | None,
) -> tuple[int, tuple[PreparedPage, ...]]:
    try:
        document = pdfium.PdfDocument(path)
    except Exception as exc:
        raise DocumentError(
            "The PDF could not be opened. Password-protected or damaged PDFs are not supported."
        ) from exc

    try:
        total_pages = len(document)
        if total_pages < 1:
            raise DocumentError("The PDF does not contain any pages.")

        page_count = min(total_pages, MAX_PAGES)
        pages = []
        for index in range(page_count):
            if on_progress:
                on_progress(index + 1, page_count)
            pages.append(_render_pdf_page(document, index))
        return total_pages, tuple(pages)
    finally:
        document.close()


def prepare_document(
    file_path: str | os.PathLike[str],
    *,
    delete_source: bool = True,
    enforce_temp_location: bool = True,
    on_progress: Callable[[int, int], None] | None = None,
) -> PreparedDocument:
    path = Path(file_path)
    if not path.is_file():
        raise DocumentError("Choose a document to continue.")
    if enforce_temp_location:
        _validate_upload_location(path)

    size = path.stat().st_size
    if size <= 0:
        raise DocumentError("The uploaded file is empty.")
    if size > MAX_DOCUMENT_BYTES:
        raise DocumentError("Files must be 20 MB or smaller.")

    name = path.name
    try:
        kind = detect_document_kind(path)
        if kind == "pdf":
            total_pages, pages = _prepare_pdf(path, on_progress)
        else:
            pages = _prepare_image(path)
            total_pages = 1

        return PreparedDocument(
            name=name,
            size=size,
            kind=kind,
            total_pages=total_pages,
            pages=pages,
        )
    finally:
        if delete_source:
            try:
                path.unlink(missing_ok=True)
            except OSError:
                pass