Estacion-de-Documentos/backend/app/services/processors/pdf_processor.py

from __future__ import annotations

from pathlib import Path

import fitz

from backend.app.config import settings
from backend.app.services.markdown_builder import title_from_path
from backend.app.services.processors.base import ProcessedContent
from backend.app.services.errors import ProcessingError
from backend.app.services.providers.mistral_client import MistralClient


class PdfProcessor:
    def __init__(self, use_ocr: bool = False, page_ranges: str | None = None):
        self.use_ocr = use_ocr
        self.page_ranges = page_ranges

    def process(self, path: Path) -> ProcessedContent:
        if self.use_ocr:
            return self._process_ocr(path)
        doc = fitz.open(path)
        pages: list[str] = []
        for index in self._selected_pages(len(doc)):
            page = doc[index - 1]
            text = page.get_text("text").strip()
            if text:
                pages.append(f"## Pagina {index}\n\n{text}")
        body = "\n\n".join(pages).strip()
        if len(body) < 80:
            raise RuntimeError("El PDF tiene poco texto extraible; procesalo con OCR")
        return ProcessedContent(title=title_from_path(path), body=body, processor="pdf-text")

    def _process_ocr(self, path: Path) -> ProcessedContent:
        try:
            doc = fitz.open(path)
        except Exception as exc:
            raise ProcessingError("BAD_PDF", f"No se pudo abrir el PDF: {exc}") from exc
        selected = self._selected_pages(len(doc))
        client = MistralClient()
        page_images: dict[int, Path] = {}
        try:
            for index in selected:
                page = doc[index - 1]
                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
                image_path = path.with_name(f"{path.stem}-page-{index}.jpg")
                pix.save(image_path)
                page_images[index] = image_path

            results = self._ocr_pages(client, page_images)
        finally:
            for image_path in page_images.values():
                image_path.unlink(missing_ok=True)

        parts: list[str] = []
        for index in selected:
            markdown = results.get(index, "")
            parts.append(f"## Pagina {index}\n\n{markdown}")
        return ProcessedContent(
            title=title_from_path(path),
            body="\n\n".join(parts),
            processor="mistral-ocr",
        )

    def _ocr_pages(self, client: MistralClient, page_images: dict[int, Path]) -> dict[int, str]:
        """OCR de las paginas con lote + fallback sincrono.

        Estrategia:
        1. Si hay >= ``batch_threshold`` paginas, intenta el job en lote.
        2. Cualquier pagina no resuelta (o si el lote entero falla) cae a
           OCR sincrono pagina a pagina con reintentos.
        3. Las paginas que fallen de forma definitiva quedan como un
           placeholder breve con el codigo de error rastreable.
        """
        indices = list(page_images.keys())
        results: dict[int, str] = {}
        page_errors: dict[int, str] = {}

        if len(indices) >= settings.mistral_ocr_batch_threshold:
            ordered = sorted(indices)
            paths = [page_images[i] for i in ordered]
            try:
                batch_results = client.ocr_images_batch(paths)
                for pos, index in enumerate(ordered):
                    if pos in batch_results:
                        results[index] = batch_results[pos]
            except ProcessingError:
                # El lote fallo; continua al fallback sincrono para todas
                # las paginas no resueltas.
                pass

        for index in indices:
            if index in results:
                continue
            try:
                results[index] = client.ocr_image(page_images[index])
            except ProcessingError as exc:
                page_errors[index] = exc.labeled
            except Exception as exc:
                page_errors[index] = f"[OCR_ERROR] {exc}"

        for index in indices:
            if index not in results:
                results[index] = f"> ⚠️ No se pudo procesar esta pagina: {page_errors.get(index, '[UNKNOWN]')}"
        return results

    def _selected_pages(self, total_pages: int) -> list[int]:
        if total_pages < 1:
            raise RuntimeError("El PDF no tiene paginas")
        if not self.page_ranges or not self.page_ranges.strip():
            return list(range(1, total_pages + 1))

        selected: list[int] = []
        for raw_part in self.page_ranges.split(","):
            part = raw_part.strip()
            if not part:
                continue
            if "-" in part:
                start_text, end_text = [piece.strip() for piece in part.split("-", 1)]
                if not start_text.isdigit() or not end_text.isdigit():
                    raise RuntimeError(f"Rango de paginas invalido: {part}")
                start = int(start_text)
                end = int(end_text)
                if start > end:
                    raise RuntimeError(f"Rango de paginas invertido: {part}")
                selected.extend(range(start, end + 1))
            else:
                if not part.isdigit():
                    raise RuntimeError(f"Pagina invalida: {part}")
                selected.append(int(part))

        unique_pages = sorted(set(selected))
        invalid = [page for page in unique_pages if page < 1 or page > total_pages]
        if invalid:
            raise RuntimeError(f"Paginas fuera de rango: {', '.join(map(str, invalid))}. El PDF tiene {total_pages} paginas")
        if not unique_pages:
            raise RuntimeError("No se seleccionaron paginas para procesar")
        return unique_pages