processors/json_formatter.py

"""
Formatador de saída JSON.

Este módulo contém funções e classes para formatar documentos
processados em formato JSON estruturado.
"""

import json
from datetime import datetime
from pathlib import Path
from typing import Any

from utils.logger import get_logger

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


def format_to_json(
    processed_data: dict[str, Any],
    filename: str,
    include_raw_content: bool = True,
    pretty_print: bool = True
) -> str:
    """
    Formata dados processados em JSON estruturado.

    Args:
        processed_data: Dados retornados pelo DoclingProcessor.
        filename: Nome do arquivo original.
        include_raw_content: Se deve incluir conteúdo completo.
        pretty_print: Se deve formatar com indentação.

    Returns:
        String JSON formatada.
    """
    document = processed_data.get("document")
    metadata = processed_data.get("metadata", {})
    tables = processed_data.get("tables", [])
    language = processed_data.get("language", "desconhecido")

    # Estrutura de saída
    output = {
        "arquivo": filename,
        "idioma": language,
        "processado_em": datetime.now().isoformat(),
        "metadados": metadata,
        "tabelas": tables,
    }

    # Adiciona conteúdo
    if include_raw_content and document:
        try:
            # Tenta exportar para dict
            if hasattr(document, "export_to_dict"):
                output["conteudo"] = document.export_to_dict()
            elif hasattr(document, "export_to_markdown"):
                output["conteudo_markdown"] = document.export_to_markdown()
            elif hasattr(document, "export_to_text"):
                output["conteudo_texto"] = document.export_to_text()
        except Exception as e:
            logger.warning(f"Erro ao exportar conteúdo: {e}")
            output["conteudo"] = None
            output["erro_exportacao"] = str(e)

    # Adiciona tempo de processamento se disponível
    if "processing_time_seconds" in processed_data:
        output["tempo_processamento_segundos"] = processed_data["processing_time_seconds"]

    # Serializa para JSON
    indent = 2 if pretty_print else None

    try:
        return json.dumps(
            output,
            ensure_ascii=False,
            indent=indent,
            default=_json_serializer
        )
    except Exception as e:
        logger.error(f"Erro ao serializar JSON: {e}")
        # Fallback: tenta sem conteúdo complexo
        output.pop("conteudo", None)
        output["erro_serializacao"] = str(e)
        return json.dumps(output, ensure_ascii=False, indent=indent)


def _json_serializer(obj: Any) -> Any:
    """
    Serializador customizado para objetos não-JSON.

    Args:
        obj: Objeto a serializar.

    Returns:
        Representação serializável do objeto.
    """
    if hasattr(obj, "isoformat"):
        return obj.isoformat()
    if hasattr(obj, "__dict__"):
        return obj.__dict__
    if isinstance(obj, bytes):
        return obj.decode("utf-8", errors="replace")
    if isinstance(obj, set):
        return list(obj)
    if isinstance(obj, Path):
        return str(obj)

    return str(obj)


class JSONFormatter:
    """
    Classe para formatação JSON com configurações personalizadas.

    Permite manter configurações consistentes entre múltiplas formatações.
    """

    def __init__(
        self,
        include_raw_content: bool = True,
        pretty_print: bool = True,
        include_tables: bool = True,
        include_metadata: bool = True
    ):
        """
        Inicializa o formatador JSON.

        Args:
            include_raw_content: Se deve incluir conteúdo completo.
            pretty_print: Se deve formatar com indentação.
            include_tables: Se deve incluir tabelas extraídas.
            include_metadata: Se deve incluir metadados.
        """
        self.include_raw_content = include_raw_content
        self.pretty_print = pretty_print
        self.include_tables = include_tables
        self.include_metadata = include_metadata

    def format(
        self,
        processed_data: dict[str, Any],
        filename: str
    ) -> str:
        """
        Formata dados processados em JSON.

        Args:
            processed_data: Dados do DoclingProcessor.
            filename: Nome do arquivo original.

        Returns:
            String JSON formatada.
        """
        # Copia para não modificar original
        data = processed_data.copy()

        # Remove elementos não desejados
        if not self.include_tables:
            data["tables"] = []

        if not self.include_metadata:
            data["metadata"] = {}

        return format_to_json(
            data,
            filename,
            include_raw_content=self.include_raw_content,
            pretty_print=self.pretty_print
        )

    def format_batch(
        self,
        items: list[tuple[dict[str, Any], str]]
    ) -> str:
        """
        Formata múltiplos documentos em um único JSON.

        Args:
            items: Lista de tuplas (processed_data, filename).

        Returns:
            String JSON com array de documentos.
        """
        documents = []

        for processed_data, filename in items:
            # Formata individualmente e converte de volta para dict
            json_str = self.format(processed_data, filename)
            doc = json.loads(json_str)
            documents.append(doc)

        indent = 2 if self.pretty_print else None

        return json.dumps(
            {"documentos": documents, "total": len(documents)},
            ensure_ascii=False,
            indent=indent
        )


def save_json(
    content: str | dict,
    output_path: str | Path,
    encoding: str = "utf-8"
) -> Path:
    """
    Salva conteúdo JSON em arquivo.

    Args:
        content: String JSON ou dicionário.
        output_path: Caminho do arquivo de saída.
        encoding: Encoding do arquivo.

    Returns:
        Path para o arquivo salvo.
    """
    output_path = Path(output_path)

    if isinstance(content, dict):
        content = json.dumps(content, ensure_ascii=False, indent=2)

    output_path.write_text(content, encoding=encoding)
    logger.debug(f"JSON salvo: {output_path}")

    return output_path