document_processor.py

#!/usr/bin/env python3
"""
Document Processor for uploaded files
Supports PDF, TXT, MD, and other text files
"""

import os
import logging
from pathlib import Path
from typing import Dict, List, Optional
import json

logger = logging.getLogger(__name__)

# Try to import PDF processing libraries
try:
    import PyPDF2
    PDF_SUPPORT = True
except ImportError:
    try:
        import pdfplumber
        PDF_SUPPORT = True
        USE_PDFPLUMBER = True
    except ImportError:
        PDF_SUPPORT = False
        USE_PDFPLUMBER = False
        logger.warning("PDF support not available. Install PyPDF2 or pdfplumber for PDF support.")


def extract_text_from_file(file_path: str, file_type: Optional[str] = None) -> Dict[str, any]:
    """
    Extract text from an uploaded file.
    
    Args:
        file_path: Path to the uploaded file
        file_type: File extension (e.g., 'pdf', 'txt'). If None, inferred from filename.
    
    Returns:
        Dict with 'text', 'metadata', and 'success' fields
    """
    file_path = Path(file_path)
    
    if not file_path.exists():
        return {
            'success': False,
            'error': f'File not found: {file_path}',
            'text': '',
            'metadata': {}
        }
    
    # Infer file type if not provided
    if not file_type:
        file_type = file_path.suffix.lower().lstrip('.')
    
    metadata = {
        'filename': file_path.name,
        'file_type': file_type,
        'file_size': file_path.stat().st_size
    }
    
    try:
        if file_type == 'pdf':
            if not PDF_SUPPORT:
                return {
                    'success': False,
                    'error': 'PDF support not available. Install PyPDF2 or pdfplumber.',
                    'text': '',
                    'metadata': metadata
                }
            
            text = extract_text_from_pdf(file_path)
            metadata['pages'] = len(text.split('\n\n')) if text else 0
            
        elif file_type in ['txt', 'text', 'md', 'markdown', 'py', 'js', 'html', 'css', 'json']:
            # Text-based files
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                text = f.read()
            metadata['lines'] = len(text.split('\n'))
            
        else:
            # Try to read as text anyway
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    text = f.read()
                metadata['lines'] = len(text.split('\n'))
            except Exception as e:
                return {
                    'success': False,
                    'error': f'Unsupported file type: {file_type}. Error: {str(e)}',
                    'text': '',
                    'metadata': metadata
                }
        
        # Clean up text
        text = text.strip()
        
        if not text:
            return {
                'success': False,
                'error': 'File appears to be empty or contains no extractable text',
                'text': '',
                'metadata': metadata
            }
        
        metadata['text_length'] = len(text)
        metadata['word_count'] = len(text.split())
        
        return {
            'success': True,
            'text': text,
            'metadata': metadata
        }
        
    except Exception as e:
        logger.error(f"Error extracting text from {file_path}: {e}", exc_info=True)
        return {
            'success': False,
            'error': f'Error processing file: {str(e)}',
            'text': '',
            'metadata': metadata
        }


def extract_text_from_pdf(file_path: Path) -> str:
    """Extract text from PDF file"""
    text_parts = []
    
    if PDF_SUPPORT:
        if USE_PDFPLUMBER:
            # Use pdfplumber (better for complex PDFs)
            try:
                import pdfplumber
                with pdfplumber.open(file_path) as pdf:
                    for page in pdf.pages:
                        page_text = page.extract_text()
                        if page_text:
                            text_parts.append(page_text)
            except Exception as e:
                logger.warning(f"pdfplumber failed, trying PyPDF2: {e}")
                # Fall back to PyPDF2
                import PyPDF2
                with open(file_path, 'rb') as f:
                    pdf_reader = PyPDF2.PdfReader(f)
                    for page in pdf_reader.pages:
                        text_parts.append(page.extract_text())
        else:
            # Use PyPDF2
            import PyPDF2
            with open(file_path, 'rb') as f:
                pdf_reader = PyPDF2.PdfReader(f)
                for page in pdf_reader.pages:
                    text_parts.append(page.extract_text())
    
    return '\n\n'.join(text_parts)


def save_document(document_data: Dict, documents_dir: Path) -> Dict:
    """
    Save processed document to storage.
    
    Args:
        document_data: Dict with 'text', 'metadata', etc.
        documents_dir: Directory to save documents
    
    Returns:
        Dict with saved document info including 'document_id'
    """
    documents_dir.mkdir(exist_ok=True)
    
    import uuid
    document_id = str(uuid.uuid4())
    filename = document_data['metadata']['filename']
    
    # Save document JSON
    doc_file = documents_dir / f"{document_id}.json"
    document_info = {
        'document_id': document_id,
        'filename': filename,
        'text': document_data['text'],
        'metadata': document_data['metadata'],
        'uploaded_at': document_data.get('uploaded_at'),
        'file_path': str(doc_file)
    }
    
    with open(doc_file, 'w', encoding='utf-8') as f:
        json.dump(document_info, f, indent=2, ensure_ascii=False)
    
    logger.info(f"Saved document {document_id}: {filename}")
    
    return document_info


def load_document(document_id: str, documents_dir: Path) -> Optional[Dict]:
    """Load a document by ID"""
    doc_file = documents_dir / f"{document_id}.json"
    if not doc_file.exists():
        return None
    
    try:
        with open(doc_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading document {document_id}: {e}")
        return None


def list_documents(documents_dir: Path) -> List[Dict]:
    """List all uploaded documents"""
    if not documents_dir.exists():
        return []
    
    documents = []
    for doc_file in documents_dir.glob("*.json"):
        try:
            with open(doc_file, 'r', encoding='utf-8') as f:
                doc = json.load(f)
                # Return summary (not full text)
                doc_summary = {
                    'document_id': doc.get('document_id'),
                    'filename': doc.get('filename'),
                    'metadata': doc.get('metadata', {}),
                    'uploaded_at': doc.get('uploaded_at'),
                    'text_length': len(doc.get('text', '')),
                    'word_count': doc.get('metadata', {}).get('word_count', 0)
                }
                documents.append(doc_summary)
        except Exception as e:
            logger.error(f"Error loading document {doc_file}: {e}")
    
    # Sort by uploaded_at (newest first)
    documents.sort(key=lambda x: x.get('uploaded_at', ''), reverse=True)
    return documents


def delete_document(document_id: str, documents_dir: Path) -> bool:
    """Delete a document"""
    doc_file = documents_dir / f"{document_id}.json"
    if doc_file.exists():
        doc_file.unlink()
        logger.info(f"Deleted document {document_id}")
        return True
    return False