preprocessing/preprocessing_modules/modular_preprocessor.py

"""
Modular Document Preprocessor

Main orchestrator class that uses all preprocessing modules to process documents.
"""

import os
import asyncio
from typing import List, Dict, Any, Union
from pathlib import Path

from config.config import OUTPUT_DIR
from .pdf_downloader import PDFDownloader
from .file_downloader import FileDownloader
from .text_extractor import TextExtractor
from .text_chunker import TextChunker
from .embedding_manager import EmbeddingManager
from .vector_storage import VectorStorage
from .metadata_manager import MetadataManager

# Import new extractors
from .docx_extractor import extract_docx
from .pptx_extractor import extract_pptx
from .xlsx_extractor import extract_xlsx
from .image_extractor import extract_image_content


class ModularDocumentPreprocessor:
    """
    Modular document preprocessor that orchestrates the entire preprocessing pipeline.
    
    This class combines all preprocessing modules to provide a clean interface
    for document processing while maintaining separation of concerns.
    """
    
    def __init__(self):
        """Initialize the modular document preprocessor."""
        # Set up base database path
        self.base_db_path = Path(OUTPUT_DIR).resolve()
        self._ensure_base_directory()
        
        # Initialize all modules
        self.pdf_downloader = PDFDownloader()  # Keep for backward compatibility
        self.file_downloader = FileDownloader()  # New enhanced downloader
        self.text_extractor = TextExtractor()
        self.text_chunker = TextChunker()
        self.embedding_manager = EmbeddingManager()
        self.vector_storage = VectorStorage(self.base_db_path)
        self.metadata_manager = MetadataManager(self.base_db_path)
        
        print("✅ Modular Document Preprocessor initialized successfully")
    
    def _ensure_base_directory(self):
        """Ensure the base directory exists."""
        if not self.base_db_path.exists():
            try:
                self.base_db_path.mkdir(parents=True, exist_ok=True)
                print(f"✅ Created directory: {self.base_db_path}")
            except PermissionError:
                print(f"⚠️  Directory {self.base_db_path} should exist in production environment")
                if not self.base_db_path.exists():
                    raise RuntimeError(f"Required directory {self.base_db_path} does not exist and cannot be created")
    
    # Delegate metadata operations to metadata manager
    def generate_doc_id(self, document_url: str) -> str:
        """Generate a unique document ID from the URL."""
        return self.metadata_manager.generate_doc_id(document_url)
    
    def is_document_processed(self, document_url: str) -> bool:
        """Check if a document has already been processed."""
        return self.metadata_manager.is_document_processed(document_url)
    
    def get_document_info(self, document_url: str) -> Dict[str, Any]:
        """Get information about a processed document."""
        return self.metadata_manager.get_document_info(document_url)
    
    def list_processed_documents(self) -> Dict[str, Dict]:
        """List all processed documents."""
        return self.metadata_manager.list_processed_documents()
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """Get statistics about all collections."""
        return self.metadata_manager.get_collection_stats()
    
    async def process_document(self, document_url: str, force_reprocess: bool = False, timeout: int = 300) -> Union[str, List]:
        """
        Process a single document: download, extract, chunk, embed, and store.
        
        Args:
            document_url: URL of the document (PDF, DOCX, PPTX, XLSX, images, etc.)
            force_reprocess: If True, reprocess even if already processed
            timeout: Download timeout in seconds (default: 300s/5min)
            
        Returns:
            str: Document ID for normal processing
            List: [content, type] for special handling (oneshot, tabular, image)
        """
        doc_id = self.generate_doc_id(document_url)
        
        # Check if already processed
        if not force_reprocess and self.is_document_processed(document_url):
            print(f"✅ Document {doc_id} already processed, skipping...")
            return doc_id
        
        print(f"🚀 Processing document: {doc_id}")
        print(f"📄 URL: {document_url}")
        
        temp_file_path = None
        try:
            # Step 1: Download file (enhanced to handle multiple types)
            temp_file_path, ext = await self.file_downloader.download_file(document_url, timeout=timeout)
            
            if temp_file_path == 'not supported':
                return ['unsupported', ext]

            # Step 2: Extract text based on file type
            full_text = ""
            match ext:
                case 'pdf':
                    full_text = await self.text_extractor.extract_text_from_pdf(temp_file_path)
                
                case 'docx':
                    full_text = extract_docx(temp_file_path)
                
                case 'pptx':
                    full_text = extract_pptx(temp_file_path)
                    return [full_text, 'oneshot']
                
                case 'url':
                    new_context = "URL for Context: " + temp_file_path
                    return [new_context, 'oneshot']
                
                case 'txt':
                    with open(temp_file_path, 'r', encoding='utf-8') as f:
                        full_text = f.read()
                
                case 'xlsx':
                    full_text = extract_xlsx(temp_file_path)
                    # Print a short preview (10-15 chars) to verify extraction
                    try:
                        preview = ''.join(full_text.split())[:15]
                        if preview:
                            print(f"🔎 XLSX extracted preview: {preview}")
                    except Exception:
                        pass
                    return [full_text, 'tabular']
                
                case 'csv':
                    with open(temp_file_path, 'r', encoding='utf-8') as f:
                        full_text = f.read()
                    return [full_text, 'tabular']

                case 'png' | 'jpeg' | 'jpg':
                    # Don't clean up image files - they'll be cleaned up by the caller
                    return [temp_file_path, 'image', True]  # Third element indicates no cleanup needed
                
                case _:
                    raise Exception(f"Unsupported file type: {ext}")

            # Validate extracted text
            if not self.text_extractor.validate_extracted_text(full_text):
                raise Exception("No meaningful text extracted from document")
            
            # Step 3: Create chunks
            chunks = self.text_chunker.chunk_text(full_text)
            
            # Check if document is too short for chunking
            if len(chunks) < 16:
                print(f"Only {len(chunks)} chunks formed, going for oneshot.")
                return [full_text, 'oneshot']
            
            if not chunks:
                raise Exception("No chunks created from text")
            
            # Log chunk statistics
            chunk_stats = self.text_chunker.get_chunk_stats(chunks)
            print(f"📊 Chunk Statistics: {chunk_stats['total_chunks']} chunks, "
                  f"avg size: {chunk_stats['avg_chunk_size']:.0f} chars")
            
            # Step 4: Create embeddings
            embeddings = await self.embedding_manager.create_embeddings(chunks)
            
            # Validate embeddings
            if not self.embedding_manager.validate_embeddings(embeddings, len(chunks)):
                raise Exception("Invalid embeddings generated")
            
            # Step 5: Store in Qdrant
            await self.vector_storage.store_in_qdrant(chunks, embeddings, doc_id)
            
            # Step 6: Save metadata
            self.metadata_manager.save_document_metadata(chunks, doc_id, document_url)
            
            print(f"✅ Document {doc_id} processed successfully: {len(chunks)} chunks")
            return doc_id
            
        except Exception as e:
            print(f"❌ Error processing document {doc_id}: {str(e)}")
            raise
        finally:
            # Clean up temporary file - but NOT for images since they need the file path
            if temp_file_path and ext not in ['png', 'jpeg', 'jpg']:
                self.file_downloader.cleanup_temp_file(temp_file_path)
    
    async def process_multiple_documents(self, document_urls: List[str], force_reprocess: bool = False) -> Dict[str, str]:
        """
        Process multiple documents concurrently.
        
        Args:
            document_urls: List of PDF URLs
            force_reprocess: If True, reprocess even if already processed
            
        Returns:
            Dict[str, str]: Mapping of URLs to document IDs
        """
        print(f"🚀 Processing {len(document_urls)} documents...")
        
        results = {}
        
        # Process documents concurrently (with limited concurrency)
        semaphore = asyncio.Semaphore(3)  # Limit to 3 concurrent downloads
        
        async def process_single(url):
            async with semaphore:
                try:
                    doc_id = await self.process_document(url, force_reprocess)
                    return url, doc_id
                except Exception as e:
                    print(f"❌ Failed to process {url}: {str(e)}")
                    return url, None
        
        tasks = [process_single(url) for url in document_urls]
        completed_tasks = await asyncio.gather(*tasks, return_exceptions=True)
        
        for result in completed_tasks:
            if isinstance(result, tuple):
                url, doc_id = result
                if doc_id:
                    results[url] = doc_id
        
        print(f"✅ Successfully processed {len(results)}/{len(document_urls)} documents")
        return results
    
    def get_system_info(self) -> Dict[str, Any]:
        """
        Get information about the preprocessing system.
        
        Returns:
            Dict[str, Any]: System information
        """
        return {
            "base_db_path": str(self.base_db_path),
            "embedding_model": self.embedding_manager.get_model_info(),
            "text_chunker_config": {
                "chunk_size": self.text_chunker.chunk_size,
                "chunk_overlap": self.text_chunker.chunk_overlap
            },
            "processed_documents_registry": self.metadata_manager.get_registry_path(),
            "collection_stats": self.get_collection_stats()
        }
    
    def cleanup_document(self, document_url: str) -> bool:
        """
        Remove all data for a specific document.
        
        Args:
            document_url: URL of the document to clean up
            
        Returns:
            bool: True if successfully cleaned up
        """
        doc_id = self.generate_doc_id(document_url)
        
        try:
            # Remove vector storage
            vector_removed = self.vector_storage.delete_collection(doc_id)
            
            # Remove metadata
            metadata_removed = self.metadata_manager.remove_document_metadata(doc_id)
            
            success = vector_removed and metadata_removed
            if success:
                print(f"✅ Successfully cleaned up document {doc_id}")
            else:
                print(f"⚠️ Partial cleanup for document {doc_id}")
            
            return success
            
        except Exception as e:
            print(f"❌ Error cleaning up document {doc_id}: {e}")
            return False