"""Render paralelo de PDF para numpy arrays via pypdfium2."""

from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import TYPE_CHECKING, Callable

if TYPE_CHECKING:
    import numpy as np


def _render_page(args: tuple[str, int, int]) -> tuple[int, "np.ndarray"]:
    """Renderiza uma página do PDF como numpy array (roda em processo separado)."""
    import numpy as np
    import pypdfium2 as pdfium

    pdf_path, page_idx, scale = args
    pdf = pdfium.PdfDocument(pdf_path)
    page = pdf[page_idx]
    img = page.render(scale=scale).to_pil()
    return page_idx, np.array(img)


def get_page_count(pdf_path: str | Path) -> int:
    """Retorna o número de páginas de um PDF."""
    import pypdfium2 as pdfium

    pdf = pdfium.PdfDocument(str(pdf_path))
    count = len(pdf)
    del pdf
    return count


def render_all_parallel(
    pdf_path: str | Path,
    num_pages: int,
    scale: int = 2,
    workers: int = 8,
    on_progress: Callable[[int, int], None] | None = None,
) -> list["np.ndarray"]:
    """Converte todas as páginas em paralelo usando ProcessPoolExecutor.

    Args:
        pdf_path: Caminho do PDF.
        num_pages: Número de páginas.
        scale: Escala de renderização.
        workers: Número de workers paralelos.
        on_progress: Callback(done, total) chamado a cada página concluída.

    Returns:
        Lista de arrays numpy ordenados por página.
    """
    pdf_str = str(pdf_path)
    args = [(pdf_str, i, scale) for i in range(num_pages)]
    images: list[np.ndarray | None] = [None] * num_pages

    with ProcessPoolExecutor(max_workers=workers) as executor:
        futures = {executor.submit(_render_page, a): a[1] for a in args}
        done = 0
        for future in as_completed(futures):
            idx, img = future.result()
            images[idx] = img
            done += 1
            if on_progress:
                on_progress(done, num_pages)

    return images  # type: ignore[return-value]
