src/tools/document_extractor.py

import os
from typing import List, Tuple
from PyPDF2 import PdfReader
from docx import Document
import mimetypes


def extract_documents_content(document_files: List[str]) -> str:
    """
    Extrai o conteúdo de múltiplos arquivos de documentos (PDF, DOC, DOCX, TXT, etc.).
    
    Args:
        document_files: Lista de caminhos para os arquivos de documentos
        
    Returns:
        String contendo o conteúdo de todos os documentos no formato:
        [NOME_DO_ARQUIVO]
        [CONTEUDO_EXTRAIDO]
        [NOME_DO_ARQUIVO2]
        [CONTEUDO_EXTRAIDO2]
        ...
    """
    results = []
    
    for file_path in document_files:
        try:
            # Verifica se o arquivo existe
            if not os.path.exists(file_path):
                error_msg = f"ERRO: Arquivo não encontrado - {file_path}"
                results.append((os.path.basename(file_path), error_msg))
                continue
            
            # Obtém apenas o nome do arquivo (sem o caminho)
            filename = os.path.basename(file_path)
            
            # Determina o tipo de arquivo
            file_extension = os.path.splitext(file_path)[1].lower()
            
            # Extrai conteúdo baseado no tipo de arquivo
            content = extract_single_document_content(file_path, file_extension)
            
            results.append((filename, content))
                
        except Exception as e:
            error_msg = f"ERRO: {str(e)}"
            results.append((os.path.basename(file_path), error_msg))
    
    # Constrói o texto final no formato solicitado
    final_text = ""
    for filename, content in results:
        final_text += f"Nome Identificador do arquivo: {filename}\n"
        final_text += f"[{content}]\n\n"
    
    return final_text.strip()


def extract_single_document_content(file_path: str, file_extension: str = None) -> str:
    """
    Extrai o conteúdo de um único arquivo de documento.
    
    Args:
        file_path: Caminho para o arquivo
        file_extension: Extensão do arquivo (opcional, será detectada automaticamente se não fornecida)
        
    Returns:
        Conteúdo extraído do documento como string
    """
    if file_extension is None:
        file_extension = os.path.splitext(file_path)[1].lower()
    
    try:
        if file_extension == '.pdf':
            return extract_pdf_content(file_path)
        elif file_extension in ['.docx', '.doc']:
            return extract_word_content(file_path)
        elif file_extension == '.txt':
            return extract_txt_content(file_path)
        else:
            return f"ERRO: Tipo de arquivo não suportado - {file_extension}"
            
    except Exception as e:
        return f"ERRO ao extrair conteúdo: {str(e)}"


def extract_pdf_content(pdf_path: str) -> str:
    """Extrai conteúdo de arquivo PDF."""
    with open(pdf_path, 'rb') as file:
        pdf_reader = PdfReader(file)
        content = ""
        for page_num in range(len(pdf_reader.pages)):
            page = pdf_reader.pages[page_num]
            page_content = page.extract_text()
            content += page_content + "\n"
        
        # Limpa e reorganiza o conteúdo
        cleaned_content = clean_pdf_content(content)
        return cleaned_content.strip()


def clean_pdf_content(content: str) -> str:
    """
    Limpa e reorganiza o conteúdo extraído do PDF para melhor legibilidade.
    
    Args:
        content: Conteúdo bruto extraído do PDF
        
    Returns:
        Conteúdo limpo e reorganizado
    """
    if not content:
        return ""
    
    # Remove caracteres especiais problemáticos
    content = content.replace('\x00', '')  # NULL bytes
    content = content.replace('\x0b', '')  # Vertical tab
    content = content.replace('\x0c', '')  # Form feed
    
    # Divide em linhas
    lines = content.split('\n')
    cleaned_lines = []
    
    for line in lines:
        line = line.strip()
        if line:  # Remove linhas vazias
            # Remove espaços múltiplos
            line = ' '.join(line.split())
            cleaned_lines.append(line)
    
    # Reorganiza o conteúdo para melhor estrutura
    reorganized_content = []
    current_section = []
    
    for line in cleaned_lines:
        # Se a linha parece ser um cabeçalho ou título (maiúsculas)
        if line.isupper() and len(line) > 3:
            # Se temos conteúdo na seção atual, adiciona ao resultado
            if current_section:
                reorganized_content.extend(current_section)
                reorganized_content.append("")  # Linha em branco para separar
            current_section = [line]
        else:
            current_section.append(line)
    
    # Adiciona a última seção
    if current_section:
        reorganized_content.extend(current_section)
    
    # Junta tudo de volta
    result = '\n'.join(reorganized_content)
    
    # Remove linhas em branco múltiplas
    while '\n\n\n' in result:
        result = result.replace('\n\n\n', '\n\n')
    
    return result


def extract_word_content(doc_path: str) -> str:
    """Extrai conteúdo de arquivo Word (DOCX/DOC)."""
    doc = Document(doc_path)
    content = ""
    for paragraph in doc.paragraphs:
        content += paragraph.text + "\n"
    return content.strip()


def extract_txt_content(txt_path: str) -> str:
    """Extrai conteúdo de arquivo de texto."""
    with open(txt_path, 'r', encoding='utf-8') as file:
        return file.read().strip()


def extract_documents_content_list(document_files: List[str]) -> List[Tuple[str, str]]:
    """
    Versão alternativa que retorna uma lista de tuplas em vez de texto concatenado.
    
    Args:
        document_files: Lista de caminhos para os arquivos de documentos
        
    Returns:
        Lista de tuplas contendo (nome_do_arquivo, conteúdo_extraído)
    """
    results = []
    
    for file_path in document_files:
        try:
            if not os.path.exists(file_path):
                results.append((os.path.basename(file_path), f"ERRO: Arquivo não encontrado"))
                continue
            
            filename = os.path.basename(file_path)
            file_extension = os.path.splitext(file_path)[1].lower()
            content = extract_single_document_content(file_path, file_extension)
            
            results.append((filename, content))
                
        except Exception as e:
            results.append((os.path.basename(file_path), f"ERRO: {str(e)}"))
    
    return results


# Funções de compatibilidade (mantidas para não quebrar código existente)
def extract_pdfs_content(pdf_files: List[str]) -> List[Tuple[str, str]]:
    """Função de compatibilidade - agora usa a nova implementação."""
    return extract_documents_content_list(pdf_files)


def extract_single_pdf_content(pdf_path: str) -> Tuple[str, str]:
    """Função de compatibilidade - agora usa a nova implementação."""
    return extract_documents_content_list([pdf_path])[0]


# Exemplo de uso
if __name__ == "__main__":
    # Exemplo com diferentes tipos de arquivos
    document_files = [
        "caminho/para/arquivo1.pdf",
        "caminho/para/arquivo2.docx",
        "caminho/para/arquivo3.txt",
        "caminho/para/arquivo4.doc"
    ]
    
    # Extrai e retorna como texto grande
    large_text = extract_documents_content(document_files)
    print(large_text)
    
    print("\n" + "="*80 + "\n")
    
    # Ou extrai como lista de tuplas
    results = extract_documents_content_list(document_files)
    for filename, content in results:
        print(f"[{filename}]")
        print(f"[{content}]")
        print("-" * 50)