app/media.py

from __future__ import annotations

import io
from pathlib import Path

from PIL import Image

from app.config import Settings

try:
    from pillow_heif import register_heif_opener

    register_heif_opener()
except ImportError:
    pass


def _load_heif() -> None:
    try:
        from pillow_heif import register_heif_opener

        register_heif_opener()
    except ImportError:
        return


def pdf_first_page_jpeg(path: Path, max_edge: int) -> bytes:
    import pypdfium2 as pdfium

    pdf = pdfium.PdfDocument(str(path))
    try:
        page = pdf[0]
        bitmap = page.render(scale=150 / 72)
        image = bitmap.to_pil().convert("RGB")
    finally:
        pdf.close()
    return _pil_to_jpeg(image, max_edge)


def _pil_to_jpeg(image: Image.Image, max_edge: int) -> bytes:
    rgb = image.convert("RGB")
    rgb.thumbnail((max_edge, max_edge), Image.Resampling.LANCZOS)
    buf = io.BytesIO()
    rgb.save(buf, format="JPEG", quality=85, optimize=True)
    return buf.getvalue()


def to_jpeg_bytes(path: Path, settings: Settings) -> bytes | None:
    suffix = path.suffix.lower()
    if suffix == ".txt":
        return None
    if suffix == ".pdf":
        return pdf_first_page_jpeg(path, settings.jpeg_max_edge)
    if suffix in {".heic", ".heif"}:
        _load_heif()
    with Image.open(path) as image:
        return _pil_to_jpeg(image, settings.jpeg_max_edge)