ocr_workbench/documents.py

from __future__ import annotations

import re
import shutil
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import fitz
from PIL import Image, ImageOps


Image.MAX_IMAGE_PIXELS = 80_000_000

SUPPORTED_IMAGE_SUFFIXES = {
    ".png",
    ".jpg",
    ".jpeg",
    ".webp",
    ".bmp",
    ".tif",
    ".tiff",
}


@dataclass(frozen=True)
class DocumentPage:
    page_number: int
    image_path: Path


class DocumentError(ValueError):
    pass


def cleanup_stale_runs(max_age_hours: float = 12.0) -> None:
    """Remove old gateway work directories and archives from the ephemeral disk."""
    cutoff = time.time() - max(0.5, float(max_age_hours)) * 3600
    root = Path(tempfile.gettempdir())
    for pattern in ("ocr_workbench_*", "ocr_workbench_*_ocr_result.zip"):
        for path in root.glob(pattern):
            try:
                if path.stat().st_mtime >= cutoff:
                    continue
                if path.is_dir():
                    shutil.rmtree(path, ignore_errors=True)
                else:
                    path.unlink(missing_ok=True)
            except OSError:
                continue


def normalize_gradio_path(value: object) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    path = getattr(value, "name", None)
    return str(path or value)


def parse_page_selection(selection: str, total_pages: int, max_pages: int) -> list[int]:
    if total_pages < 1:
        return []
    max_pages = max(1, int(max_pages))
    text = (selection or "").strip()
    if not text:
        return list(range(min(total_pages, max_pages)))

    selected: set[int] = set()
    for token in re.split(r"\s*,\s*", text):
        if not token:
            continue
        if "-" in token:
            match = re.fullmatch(r"(\d+)\s*-\s*(\d+)", token)
            if not match:
                raise DocumentError(f"Invalid page range: {token!r}")
            start, end = map(int, match.groups())
            if start > end:
                start, end = end, start
            selected.update(range(start, end + 1))
        else:
            if not token.isdigit():
                raise DocumentError(f"Invalid page number: {token!r}")
            selected.add(int(token))

    invalid = sorted(page for page in selected if page < 1 or page > total_pages)
    if invalid:
        raise DocumentError(
            f"Page selection contains pages outside 1-{total_pages}: {invalid}"
        )

    ordered = sorted(selected)
    if len(ordered) > max_pages:
        raise DocumentError(
            f"Selected {len(ordered)} pages, but the current limit is {max_pages}."
        )
    return [page - 1 for page in ordered]


def _save_normalized_image(input_path: Path, output_path: Path) -> None:
    try:
        with Image.open(input_path) as image:
            normalized = ImageOps.exif_transpose(image).convert("RGB")
            normalized.save(output_path, format="PNG", optimize=False)
    except Exception as exc:
        raise DocumentError(f"Unable to decode image: {exc}") from exc


def prepare_document(
    source_path: str,
    *,
    dpi: int,
    selection: str,
    max_pages: int,
    max_input_mb: int,
) -> tuple[Path, list[DocumentPage]]:
    source = Path(source_path)
    if not source.exists() or not source.is_file():
        raise DocumentError("Upload an image or PDF first.")

    size_mb = source.stat().st_size / (1024 * 1024)
    if size_mb > max_input_mb:
        raise DocumentError(
            f"Input is {size_mb:.1f} MB; the configured limit is {max_input_mb} MB."
        )

    run_dir = Path(tempfile.mkdtemp(prefix="ocr_workbench_"))
    input_dir = run_dir / "input"
    pages_dir = run_dir / "source_pages"
    input_dir.mkdir(parents=True, exist_ok=True)
    pages_dir.mkdir(parents=True, exist_ok=True)
    shutil.copy2(source, input_dir / source.name)

    suffix = source.suffix.lower()
    pages: list[DocumentPage] = []
    if suffix == ".pdf":
        try:
            document = fitz.open(source)
        except Exception as exc:
            raise DocumentError(f"Unable to open PDF: {exc}") from exc
        try:
            indices = parse_page_selection(selection, document.page_count, max_pages)
            scale = max(72, min(360, int(dpi))) / 72.0
            matrix = fitz.Matrix(scale, scale)
            for index in indices:
                output = pages_dir / f"page_{index + 1:04d}.png"
                pixmap = document.load_page(index).get_pixmap(
                    matrix=matrix,
                    alpha=False,
                    colorspace=fitz.csRGB,
                )
                pixmap.save(output)
                pages.append(DocumentPage(index + 1, output))
        finally:
            document.close()
    elif suffix in SUPPORTED_IMAGE_SUFFIXES:
        output = pages_dir / "page_0001.png"
        _save_normalized_image(source, output)
        pages.append(DocumentPage(1, output))
    else:
        raise DocumentError(
            f"Unsupported file type {suffix!r}. Use PDF, PNG, JPEG, WebP, BMP, or TIFF."
        )

    if not pages:
        raise DocumentError("No pages were selected.")
    return run_dir, pages


def gallery_items(pages: Iterable[DocumentPage]) -> list[tuple[str, str]]:
    return [(str(page.image_path), f"Page {page.page_number}") for page in pages]