pdf-toolkit/core/compress.py

"""
Estratégias de compressão e otimização de PDF.

Usa pypdf (API oficial de ImageFile.replace) + Pillow para
re-amostragem de imagens. A abordagem antiga de gravar em
`xobj._data` quebrava o /Filter do stream e falhava em silêncio
na maioria dos PDFs (FlateDecode / pixels crus).
"""

from __future__ import annotations

import io
import logging
from typing import Optional, Tuple

from PIL import Image
from pypdf import PdfReader, PdfWriter

from config import QUALITY_PRESETS

logger = logging.getLogger(__name__)


def _prepare_image(
    img: Image.Image,
    quality: int = 60,
    max_dim: int = 2000,
) -> Tuple[Image.Image, dict]:
    """
    Normaliza modo de cor e redimensiona para o preset.

    Retorna (imagem_pronta, kwargs para Image.save / replace).
    """
    # JPEG (via replace→PDF) não aceita alpha/CMYK/paleta de forma confiável
    if img.mode in ("RGBA", "LA"):
        background = Image.new("RGB", img.size, (255, 255, 255))
        alpha = img.split()[-1]
        background.paste(img.convert("RGBA"), mask=alpha)
        img = background
    elif img.mode == "P":
        img = img.convert("RGBA")
        if "A" in img.getbands():
            background = Image.new("RGB", img.size, (255, 255, 255))
            background.paste(img, mask=img.split()[-1])
            img = background
        else:
            img = img.convert("RGB")
    elif img.mode == "CMYK":
        img = img.convert("RGB")
    elif img.mode not in ("RGB", "L"):
        img = img.convert("RGB")

    w, h = img.size
    if max(w, h) > max_dim:
        ratio = max_dim / max(w, h)
        new_size = (max(1, int(w * ratio)), max(1, int(h * ratio)))
        img = img.resize(new_size, Image.LANCZOS)

    save_kwargs = {"quality": quality, "optimize": True}
    return img, save_kwargs


def _downsample_image(
    img_bytes: bytes,
    quality: int = 60,
    max_dim: int = 2000,
    fmt: str = "JPEG",
) -> bytes:
    """
    Redimensiona e re-comprime bytes de imagem (helper legado / pdfa).

    Aceita apenas containers que o Pillow abra (JPEG/PNG/...).
    Pixels crus de streams FlateDecode NÃO funcionam aqui — use
    compress_pdf(), que passa por page.images (PIL Image já montada).
    """
    img = Image.open(io.BytesIO(img_bytes))
    img, save_kwargs = _prepare_image(img, quality=quality, max_dim=max_dim)

    buf = io.BytesIO()
    out_fmt = "JPEG" if fmt.upper() == "JPEG" else "PNG"
    if out_fmt == "JPEG" and img.mode not in ("RGB", "L"):
        img = img.convert("RGB")
    if out_fmt == "JPEG":
        img.save(buf, format="JPEG", **save_kwargs)
    else:
        img.save(buf, format="PNG", optimize=True)
    return buf.getvalue()


def compress_pdf(
    data: bytes,
    preset: str = "ebook",
    remove_unused: bool = True,
) -> bytes:
    """
    Comprime um PDF re-amostrando imagens embutidas.

    Args:
        data:       Conteúdo do PDF original.
        preset:     Perfil de qualidade ("screen", "ebook", "printer", "prepress").
        remove_unused: Tenta deduplicar objetos idênticos no writer.

    Retorna bytes do PDF (pode ser maior em PDFs só-texto; o engine
    decide se devolve o original nesse caso).
    """
    opts = QUALITY_PRESETS.get(preset, QUALITY_PRESETS["ebook"])
    quality = opts["quality"]
    max_dim = opts["max_dim"]

    reader = PdfReader(io.BytesIO(data))
    # clone_from é obrigatório para ImageFile.replace() funcionar
    writer = PdfWriter(clone_from=reader)

    replaced = 0
    errors = 0

    for page in writer.pages:
        try:
            images = list(page.images)
        except Exception as exc:
            logger.debug("Falha ao listar imagens da página: %s", exc)
            continue

        for img_file in images:
            # Imagens inline não têm indirect_reference — replace() recusa
            if getattr(img_file, "indirect_reference", None) is None:
                continue
            try:
                pil = img_file.image
                if pil is None:
                    continue
                prepared, save_kwargs = _prepare_image(
                    pil, quality=quality, max_dim=max_dim
                )
                img_file.replace(prepared, **save_kwargs)
                replaced += 1
            except Exception as exc:
                errors += 1
                logger.debug(
                    "Falha ao recomprimir imagem %s: %s",
                    getattr(img_file, "name", "?"),
                    exc,
                )

    if remove_unused:
        try:
            writer.compress_identical_objects()
        except Exception:
            pass

    out_buf = io.BytesIO()
    writer.write(out_buf)
    result = out_buf.getvalue()

    logger.info(
        "compress_pdf preset=%s images_replaced=%d errors=%d "
        "in=%d out=%d",
        preset,
        replaced,
        errors,
        len(data),
        len(result),
    )
    return result


def compress_pdf_with_stats(
    data: bytes,
    preset: str = "ebook",
    remove_unused: bool = True,
) -> Tuple[bytes, dict]:
    """
    Igual a compress_pdf, mas devolve estatísticas úteis para UI/API.
    """
    opts = QUALITY_PRESETS.get(preset, QUALITY_PRESETS["ebook"])
    quality = opts["quality"]
    max_dim = opts["max_dim"]

    reader = PdfReader(io.BytesIO(data))
    writer = PdfWriter(clone_from=reader)

    replaced = 0
    skipped = 0
    errors = 0

    for page in writer.pages:
        try:
            images = list(page.images)
        except Exception:
            continue

        for img_file in images:
            if getattr(img_file, "indirect_reference", None) is None:
                skipped += 1
                continue
            try:
                pil = img_file.image
                if pil is None:
                    skipped += 1
                    continue
                prepared, save_kwargs = _prepare_image(
                    pil, quality=quality, max_dim=max_dim
                )
                img_file.replace(prepared, **save_kwargs)
                replaced += 1
            except Exception:
                errors += 1

    if remove_unused:
        try:
            writer.compress_identical_objects()
        except Exception:
            pass

    out_buf = io.BytesIO()
    writer.write(out_buf)
    result = out_buf.getvalue()

    stats = {
        "images_replaced": replaced,
        "images_skipped": skipped,
        "images_errors": errors,
        "original_bytes": len(data),
        "result_bytes": len(result),
        "reduced": len(result) < len(data),
        "preset": preset,
    }
    return result, stats


def linearize_pdf(data: bytes) -> bytes:
    """
    Re-escreve o PDF (clone limpo). Não é linearização real de web-fast,
    mas reestrutura o arquivo via pypdf.
    """
    reader = PdfReader(io.BytesIO(data))
    writer = PdfWriter(clone_from=reader)
    out_buf = io.BytesIO()
    writer.write(out_buf)
    return out_buf.getvalue()