pdf-toolkit/core/operations.py

"""
Operações gerais sobre PDFs: merge, split, rotate, extração de
páginas/texto/imagens.
"""

import io
from typing import List, Optional, Tuple

from PIL import Image
from pypdf import PdfReader, PdfWriter, PageObject


# ── Merge ───────────────────────────────────────────────────────────

def merge_pdfs(files_data: List[Tuple[bytes, str]]) -> bytes:
    """
    Junta múltiplos PDFs em um único arquivo, na ordem fornecida.

    Args:
        files_data: Lista de (conteúdo_bytes, nome_arquivo).

    Retorna bytes do PDF mesclado.
    """
    writer = PdfWriter()

    for data, fname in files_data:
        reader = PdfReader(io.BytesIO(data))
        for page in reader.pages:
            writer.add_page(page)

    out_buf = io.BytesIO()
    writer.write_stream(out_buf)
    return out_buf.getvalue()


# ── Split ───────────────────────────────────────────────────────────

def split_pdf(
    data: bytes,
    page_ranges: List[Tuple[int, int]],
    filename_base: str = "pagina",
) -> List[Tuple[str, bytes]]:
    """
    Extrai intervalos de páginas de um PDF.

    Args:
        data:           Conteúdo do PDF.
        page_ranges:    Lista de (inicio, fim), 0-indexado. Ex: [(0,2), (5,7)].
        filename_base:  Base para nomes.

    Retorna lista de (nome, bytes).
    """
    reader = PdfReader(io.BytesIO(data))
    total = len(reader.pages)
    results = []

    for i, (start, end) in enumerate(page_ranges):
        writer = PdfWriter()
        writer.metadata = reader.metadata
        for idx in range(start, min(end + 1, total)):
            writer.add_page(reader.pages[idx])

        out_buf = io.BytesIO()
        writer.write_stream(out_buf)
        label = f"{filename_base}_p{start+1}-{min(end,total-1)+1}.pdf"
        results.append((label, out_buf.getvalue()))

    return results


def extract_pages(
    data: bytes,
    pages: List[int],
) -> bytes:
    """
    Extrai páginas específicas de um PDF e retorna um novo PDF
    contendo apenas essas páginas (na ordem fornecida).

    Args:
        data:   Conteúdo do PDF.
        pages:  Lista de índices de página (0-indexado).

    Retorna bytes do novo PDF.
    """
    reader = PdfReader(io.BytesIO(data))
    writer = PdfWriter()
    writer.metadata = reader.metadata
    for idx in pages:
        if idx < len(reader.pages):
            writer.add_page(reader.pages[idx])
    out_buf = io.BytesIO()
    writer.write_stream(out_buf)
    return out_buf.getvalue()


# ── Rotação ─────────────────────────────────────────────────────────

def rotate_pages(
    data: bytes,
    rotations: List[Tuple[int, int]],
) -> bytes:
    """
    Rotaciona páginas específicas de um PDF.

    Args:
        data:       Conteúdo do PDF.
        rotations:  Lista de (pagina_idx, angulo_em_graus).
                    Angulo = 90, 180, 270 (sentido anti-horário).

    Retorna bytes do PDF rotacionado.
    """
    reader = PdfReader(io.BytesIO(data))
    writer = PdfWriter()
    writer.metadata = reader.metadata

    rot_map = {p: a for p, a in rotations}

    for i, page in enumerate(reader.pages):
        if i in rot_map:
            page.rotate(rot_map[i])
        writer.add_page(page)

    out_buf = io.BytesIO()
    writer.write_stream(out_buf)
    return out_buf.getvalue()


# ── Extração de texto ──────────────────────────────────────────────

def extract_text(data: bytes) -> str:
    """Extrai texto de todas as páginas do PDF."""
    reader = PdfReader(io.BytesIO(data))
    pages_text = []
    for page in reader.pages:
        text = page.extract_text()
        if text:
            pages_text.append(text)
    return "\n".join(pages_text)


# ── Extração de imagens ────────────────────────────────────────────

def extract_images(data: bytes) -> List[Tuple[str, bytes]]:
    """
    Extrai todas as imagens de um PDF.

    Retorna lista de (nome_extraido, bytes_imagem).
    """
    reader = PdfReader(io.BytesIO(data))
    images = []
    counter = 0

    for page_num, page in enumerate(reader.pages):
        if "/XObject" not in page.get("/Resources", {}):
            continue
        xobjects = page["/Resources"]["/XObject"]
        for obj_name, xobj in xobjects.items():
            if xobj.get("/Subtype") == "/Image":
                try:
                    raw = xobj.get_data()
                    fmt = xobj.get("/Filter", "")
                    ext = ".jpg" if "/DCTDecode" in str(fmt) else ".png"
                    name = f"img_p{page_num+1}_{obj_name.strip('/')}{ext}"
                    images.append((name, raw))
                    counter += 1
                except Exception:
                    pass

    return images


# ── Preview de páginas (como imagens) ──────────────────────────────

def preview_pages(
    data: bytes,
    max_pages: int = 10,
    dpi: int = 100,
) -> List[Image.Image]:
    """
    Gera miniaturas/thumbnails das primeiras páginas de um PDF.

    NOTA: pypdf não renderiza páginas; este método retorna uma
    imagem em branco com o número da página e o tamanho.
    Para preview real, seria necessário pdf2image (poppler).

    Esta implementação serve como placeholder — o Streamlit
    usará pypdf para mostrar metadados e contagem de páginas,
    enquanto a renderização real depende de poppler-utils.
    """
    reader = PdfReader(io.BytesIO(data))
    previews = []
    for i, page in enumerate(reader.pages[:max_pages]):
        # Cria uma imagem placeholder com dimensões proporcionais
        w = int(page.mediabox.width) if page.mediabox else 595
        h = int(page.mediabox.height) if page.mediabox else 842
        scale = min(400 / w, 500 / h)
        img = Image.new("RGB", (int(w * scale), int(h * scale)), (245, 245, 245))
        previews.append(img)
    return previews