IPSExtractor/services/extraction/core/pdf_utils.py

"""Utilidades para convertir documentos (PDF/TIFF) a imágenes base64."""
import base64
import logging
import os
from pathlib import Path
from io import BytesIO
from typing import List, Tuple, Optional
import time

from pdf2image import convert_from_path
from PIL import Image, ImageSequence

Image.MAX_IMAGE_PIXELS = None

logger = logging.getLogger(__name__)

_SERVICE_ROOT = Path(__file__).resolve().parents[1]
_BUNDLED_WINDOWS_POPPLER = _SERVICE_ROOT / "dependencies" / "poppler-24.08.0" / "Library" / "bin"
_LEGACY_WINDOWS_POPPLER = Path("dependencies/poppler-24.08.0/Library/bin")
_DOCKER_POPPLER = Path("/usr/bin")


def _resolve_poppler(poppler_path: str | None) -> str | None:
  """Devuelve una ruta de Poppler solo si existe; si no, usa el PATH del sistema."""
  candidates = [
    Path(poppler_path.strip()) if poppler_path and poppler_path.strip() else None,
    _BUNDLED_WINDOWS_POPPLER,
    _LEGACY_WINDOWS_POPPLER,
    _DOCKER_POPPLER,
  ]
  for candidate in candidates:
    if candidate and candidate.exists():
      return str(candidate)
  return None


def pdf_to_images_bs64(
  pdf_path: str,
  validate: bool = True,
  dpi: int = 800,
  thread_count: int = 10,
  poppler_path: str | None = None,
) -> List[str]:
  """
  Convierte un archivo PDF completo a una lista de imágenes PNG en base64.

  Args:
      pdf_path: Ruta al archivo PDF.
      validate: Si es True, intenta abrir cada imagen generada para verificar su integridad.
      dpi: Resolución de la imagen (puntos por pulgada).
      thread_count: Número de hilos para la conversión paralela de páginas.
      poppler_path: Ruta personalizada a los binarios de Poppler (opcional).

  Returns:
      Lista de strings con las imágenes codificadas en base64.

  Raises:
      ValueError: Si la validación de alguna imagen falla.
  """
  poppler_bin = _resolve_poppler(poppler_path)
  t0 = time.time()
  logger.info("Convirtiendo PDF a imágenes", extra={"pdf": pdf_path, "dpi": dpi})
  images = convert_from_path(pdf_path, dpi=dpi, thread_count=thread_count, poppler_path=poppler_bin)
  resultados: List[str] = []
  for indice, imagen in enumerate(images):
    buffer = BytesIO()
    imagen.save(buffer, format="PNG")
    b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
    if validate:
      try:
        raw = base64.b64decode(b64, validate=True)
        with Image.open(BytesIO(raw)) as im:
          im.verify()
      except Exception as error:
        raise ValueError(f"Imagen base64 inválida en la página {indice + 1}: {error}")
    resultados.append(b64)
  elapsed_ms = int((time.time() - t0) * 1000)
  logger.debug(f"PDF->imgs OK (pages={len(resultados)}, ms={elapsed_ms}, threads={thread_count})", extra={"pdf": pdf_path, "dpi": dpi})
  logger.info("PDF convertido", extra={"pdf": pdf_path, "dpi": dpi})
  return resultados


def tiff_to_images_bs64(tiff_path: str, validate: bool = True, pages_range: Optional[Tuple[int, int]] = (2, 3)) -> List[str]:
  """
  Convierte un archivo TIFF multipágina a una lista de imágenes PNG en base64.

  Args:
      tiff_path: Ruta al archivo TIFF.
      validate: Si es True, verifica la integridad de las imágenes generadas.
      pages_range: Tupla (inicio, fin) de páginas a procesar (1-indexed, inclusivo).
                  Use None para procesar todas las páginas.

  Returns:
      Lista de imágenes en base64.
  """
  resultados: List[str] = []
  t0 = time.time()
  logger.info("Convirtiendo TIFF a imágenes", extra={"pdf": tiff_path})
  with Image.open(tiff_path) as im:
    frames = []
    for idx, frame in enumerate(ImageSequence.Iterator(im), start=1):
      if pages_range is not None:
        start, end = pages_range
        # Permitir extremos None para significar “sin límite” por ese lado
        if start is not None and idx < start:
          continue
        if end is not None and idx > end:
          continue
      frames.append(frame.copy())
    for indice, frame in enumerate(frames):
      buffer = BytesIO()
      frame.convert("RGB").save(buffer, format="PNG")
      b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
      if validate:
        try:
          raw = base64.b64decode(b64, validate=True)
          with Image.open(BytesIO(raw)) as imv:
            imv.verify()
        except Exception as error:
          raise ValueError(f"Imagen base64 inválida en la página TIFF {indice + 1}: {error}")
      resultados.append(b64)
  elapsed_ms = int((time.time() - t0) * 1000)
  logger.debug(f"TIFF->imgs OK (pages={len(resultados)}, ms={elapsed_ms})", extra={"pdf": tiff_path})
  logger.info("TIFF convertido", extra={"pdf": tiff_path})
  return resultados

def document_to_images_bs64(
  path: str,
  validate: bool = True,
  dpi: int = 600,
  thread_count: int = 10,
  tiff_pages: Optional[Tuple[int, int]] = (2, 3),
  tiff_page_range: Optional[Tuple[int, int]] = None,
  poppler_path: str | None = None,
) -> List[str]:
  """
  Función unificada para convertir documentos (PDF o TIFF) a imágenes base64.

  Detecta automáticamente el formato por la extensión del archivo.

  Args:
      path: Ruta al archivo.
      validate: Si es True, valida las imágenes generadas.
      dpi: Resolución para archivos PDF.
      thread_count: Hilos paralelos para archivos PDF.
      tiff_pages: Páginas TIFF por defecto si tiff_page_range es None.
      tiff_page_range: Rango explícito de páginas para TIFF (inicio, fin).
      poppler_path: Ruta personalizada a Poppler para PDFs.

  Returns:
      Lista de imágenes en base64.
  """
  ext = os.path.splitext(path)[1].lower()
  if ext == ".pdf":
    return pdf_to_images_bs64(path, validate=validate, dpi=dpi, thread_count=thread_count, poppler_path=poppler_path)
  if ext in (".tif", ".tiff"):
    effective_range = tiff_pages
    if tiff_page_range is not None:
      effective_range = None if tiff_page_range == (None, None) else tiff_page_range
    return tiff_to_images_bs64(path, validate=validate, pages_range=effective_range)
  raise ValueError(f"Formato no soportado: {ext}")