app/utils/document_extractor.py

"""
Document text extraction utilities for various file formats.
"""
import os
from typing import Optional
import PyPDF2
import docx
from unstructured.partition.auto import partition


class DocumentExtractor:
    """Extract text content from various document formats."""
    
    @staticmethod
    def extract_text(file_path: str, file_type: str) -> str:
        """
        Extract text from a document file.
        
        Args:
            file_path: Path to the document file
            file_type: File extension (pdf, txt, docx)
            
        Returns:
            Extracted text content
            
        Raises:
            ValueError: If file type is not supported
            Exception: If extraction fails
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        file_type = file_type.lower()
        
        try:
            if file_type == 'txt':
                return DocumentExtractor._extract_txt(file_path)
            elif file_type == 'pdf':
                return DocumentExtractor._extract_pdf(file_path)
            elif file_type == 'docx':
                return DocumentExtractor._extract_docx(file_path)
            else:
                raise ValueError(f"Unsupported file type: {file_type}")
                
        except Exception as e:
            print(f"Error extracting text from {file_path}: {e}")
            raise
    
    @staticmethod
    def _extract_txt(file_path: str) -> str:
        """Extract text from TXT file."""
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            return f.read()
    
    @staticmethod
    def _extract_pdf(file_path: str) -> str:
        """Extract text from PDF file using PyPDF2."""
        text_content = []
        
        try:
            with open(file_path, 'rb') as f:
                pdf_reader = PyPDF2.PdfReader(f)
                
                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    text = page.extract_text()
                    if text.strip():
                        text_content.append(text)
            
            return "\n\n".join(text_content)
            
        except Exception as e:
            print(f"PyPDF2 extraction failed, trying unstructured library: {e}")
            # Fallback to unstructured library
            return DocumentExtractor._extract_with_unstructured(file_path)
    
    @staticmethod
    def _extract_docx(file_path: str) -> str:
        """Extract text from DOCX file."""
        try:
            doc = docx.Document(file_path)
            text_content = []
            
            # Extract paragraphs
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text_content.append(paragraph.text)
            
            # Extract tables
            for table in doc.tables:
                for row in table.rows:
                    row_text = " | ".join(cell.text.strip() for cell in row.cells)
                    if row_text.strip():
                        text_content.append(row_text)
            
            return "\n\n".join(text_content)
            
        except Exception as e:
            print(f"python-docx extraction failed: {e}")
            raise
    
    @staticmethod
    def _extract_with_unstructured(file_path: str) -> str:
        """
        Extract text using unstructured library (fallback method).
        This handles complex PDFs with tables and images better.
        """
        try:
            elements = partition(filename=file_path)
            text_content = [str(element) for element in elements]
            return "\n\n".join(text_content)
        except Exception as e:
            print(f"Unstructured extraction failed: {e}")
            raise


# Global extractor instance
document_extractor = DocumentExtractor()