processors/docling_processor.py

"""
Processador principal usando Docling.

Este módulo contém a classe DoclingProcessor que é responsável por
converter documentos usando a biblioteca Docling.
"""

import time
from pathlib import Path
from typing import Any, Optional

from docling.document_converter import DocumentConverter
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
    PdfPipelineOptions,
    TableFormerMode,
)
from docling.document_converter import PdfFormatOption

from utils.logger import get_logger, ProcessingLogger

# Logger para este módulo
logger = get_logger(__name__)


class DoclingProcessor:
    """
    Processador de documentos usando Docling.

    Esta classe encapsula a lógica de conversão de documentos,
    incluindo configuração de pipeline e extração de metadados.
    """

    def __init__(
        self,
        enable_ocr: bool = True,
        enable_table_detection: bool = True,
        use_gpu: bool = True
    ):
        """
        Inicializa o processador Docling.

        Args:
            enable_ocr: Se deve habilitar OCR para imagens.
            enable_table_detection: Se deve detectar tabelas.
            use_gpu: Flag indicando se GPU está disponível.
                     Nota: A aceleração GPU real é obtida via `@spaces.GPU`
                     no app.py. Esta flag serve para logging/debug.
        """
        self.enable_ocr = enable_ocr
        self.enable_table_detection = enable_table_detection
        self.use_gpu = use_gpu

        # Configuração do pipeline
        self._setup_converter()

        logger.info(
            f"DoclingProcessor inicializado "
            f"(OCR={enable_ocr}, tabelas={enable_table_detection}, GPU={use_gpu})"
        )

    def _setup_converter(self) -> None:
        """Configura o DocumentConverter com as opções adequadas."""
        # Configurações específicas para PDF
        pipeline_options = PdfPipelineOptions()
        pipeline_options.do_ocr = self.enable_ocr
        pipeline_options.do_table_structure = self.enable_table_detection

        if self.enable_table_detection:
            # Usa TableFormer para melhor detecção de tabelas
            pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE

        # Cria o converter com as opções
        self.converter = DocumentConverter(
            format_options={
                InputFormat.PDF: PdfFormatOption(
                    pipeline_options=pipeline_options
                )
            }
        )

    def process_document(self, file_path: str | Path) -> dict[str, Any]:
        """
        Processa um documento e retorna dados estruturados.

        Args:
            file_path: Caminho para o arquivo a processar.

        Returns:
            Dicionário com documento convertido, metadados e tabelas.

        Raises:
            Exception: Se o processamento falhar.
        """
        file_path = Path(file_path)

        with ProcessingLogger(logger, "Conversão Docling", file_path.name):
            start_time = time.time()

            try:
                # Converte o documento
                result = self.converter.convert(str(file_path))

                # Extrai informações
                document = result.document

                processing_time = time.time() - start_time

                return {
                    "document": document,
                    "metadata": self._extract_metadata(result, file_path),
                    "tables": self._extract_tables(document),
                    "language": self._detect_language(document),
                    "processing_time_seconds": processing_time,
                }

            except Exception as e:
                logger.error(f"Erro ao processar {file_path.name}: {e}")
                raise

    def _extract_metadata(
        self,
        result: Any,
        file_path: Path
    ) -> dict[str, Any]:
        """
        Extrai metadados do documento processado.

        Args:
            result: Resultado da conversão Docling.
            file_path: Caminho do arquivo original.

        Returns:
            Dicionário com metadados do documento.
        """
        document = result.document

        metadata = {
            "nome_arquivo": file_path.name,
            "extensao": file_path.suffix.lower(),
            "tamanho_bytes": file_path.stat().st_size if file_path.exists() else 0,
        }

        # Tenta extrair metadados do documento
        try:
            if hasattr(document, "metadata") and document.metadata:
                doc_meta = document.metadata
                if hasattr(doc_meta, "title") and doc_meta.title:
                    metadata["titulo"] = doc_meta.title
                if hasattr(doc_meta, "author") and doc_meta.author:
                    metadata["autor"] = doc_meta.author
                if hasattr(doc_meta, "creation_date") and doc_meta.creation_date:
                    metadata["data_criacao"] = str(doc_meta.creation_date)
        except Exception as e:
            logger.debug(f"Não foi possível extrair metadados: {e}")

        # Contagem de elementos
        try:
            if hasattr(document, "pages"):
                metadata["num_paginas"] = len(list(document.pages))
            if hasattr(document, "tables"):
                metadata["num_tabelas"] = len(list(document.tables))
            if hasattr(document, "pictures"):
                metadata["num_imagens"] = len(list(document.pictures))
        except Exception as e:
            logger.debug(f"Erro ao contar elementos: {e}")

        return metadata

    def _extract_tables(self, document: Any) -> list[dict[str, Any]]:
        """
        Extrai tabelas do documento.

        Args:
            document: Documento Docling convertido.

        Returns:
            Lista de dicionários representando tabelas.
        """
        tables = []

        try:
            if not hasattr(document, "tables"):
                return tables

            for i, table in enumerate(document.tables):
                table_data = {
                    "indice": i + 1,
                    "dados": None,
                    "linhas": 0,
                    "colunas": 0,
                }

                # Tenta extrair dados da tabela
                try:
                    if hasattr(table, "export_to_dataframe"):
                        df = table.export_to_dataframe()
                        table_data["dados"] = df.to_dict(orient="records")
                        table_data["linhas"] = len(df)
                        table_data["colunas"] = len(df.columns)
                        table_data["colunas_nomes"] = list(df.columns)
                    elif hasattr(table, "to_markdown"):
                        table_data["markdown"] = table.to_markdown()
                except Exception as e:
                    logger.debug(f"Erro ao exportar tabela {i+1}: {e}")
                    # Fallback: tenta obter texto
                    if hasattr(table, "text"):
                        table_data["texto"] = table.text

                tables.append(table_data)

        except Exception as e:
            logger.warning(f"Erro ao extrair tabelas: {e}")

        logger.debug(f"Extraídas {len(tables)} tabelas")
        return tables

    def _detect_language(self, document: Any) -> str:
        """
        Detecta o idioma do documento.

        Args:
            document: Documento Docling convertido.

        Returns:
            Código do idioma detectado (ex: "pt", "en").
        """
        try:
            # Tenta usar langdetect
            from langdetect import detect, LangDetectException

            # Extrai texto do documento
            if hasattr(document, "export_to_text"):
                text = document.export_to_text()
            elif hasattr(document, "export_to_markdown"):
                text = document.export_to_markdown()
            else:
                return "desconhecido"

            # Usa apenas os primeiros 1000 caracteres para detecção
            sample = text[:1000] if text else ""

            if len(sample) < 20:
                return "desconhecido"

            lang = detect(sample)
            logger.debug(f"Idioma detectado: {lang}")
            return lang

        except LangDetectException:
            return "desconhecido"
        except ImportError:
            logger.warning("langdetect não disponível")
            return "nao_detectado"
        except Exception as e:
            logger.debug(f"Erro na detecção de idioma: {e}")
            return "erro"

    def get_markdown(self, processed_data: dict[str, Any]) -> str:
        """
        Obtém o documento em formato Markdown.

        Args:
            processed_data: Dados retornados por process_document().

        Returns:
            String com o documento em Markdown.
        """
        document = processed_data.get("document")

        if document and hasattr(document, "export_to_markdown"):
            return document.export_to_markdown()

        return ""

    def get_text(self, processed_data: dict[str, Any]) -> str:
        """
        Obtém o documento em texto puro.

        Args:
            processed_data: Dados retornados por process_document().

        Returns:
            String com o texto do documento.
        """
        document = processed_data.get("document")

        if document and hasattr(document, "export_to_text"):
            return document.export_to_text()

        return ""