data_loader.py

import os
import logging
from typing import List, Dict
from pdf_processor import PDFProcessor
from vector_store import VectorStore
from config import Config

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DataLoader:
    def __init__(self):
        self.config = Config()
        self.pdf_processor = PDFProcessor()
        self.vector_store = VectorStore()
    
    def load_documents(self, force_reload: bool = False) -> Dict:
        """Load all PDF documents from the configured directory into the vector store."""
        try:
            logger.info("Starting document loading process...")
            
            # Check if PDF directory exists
            if not os.path.exists(self.config.PDF_DIRECTORY):
                logger.error(f"PDF directory {self.config.PDF_DIRECTORY} does not exist")
                return {
                    'success': False,
                    'error': f"PDF directory {self.config.PDF_DIRECTORY} does not exist",
                    'documents_processed': 0,
                    'chunks_created': 0
                }
            
            # Get collection info
            collection_info = self.vector_store.get_collection_info()
            
            # If force reload is requested, clear the collection
            if force_reload:
                logger.info("Force reload requested. Clearing existing collection...")
                self.vector_store.clear_collection()
                collection_info = self.vector_store.get_collection_info()
            
            # Check if documents are already loaded
            if collection_info.get('document_count', 0) > 0 and not force_reload:
                logger.info(f"Documents already loaded ({collection_info['document_count']} chunks found)")
                return {
                    'success': True,
                    'message': f"Documents already loaded ({collection_info['document_count']} chunks)",
                    'documents_processed': 0,
                    'chunks_created': collection_info['document_count'],
                    'reused_existing': True
                }
            
            # Process PDFs
            logger.info("Processing PDF documents...")
            chunks = self.pdf_processor.process_pdf_directory()
            
            if not chunks:
                logger.warning("No chunks created from PDFs")
                return {
                    'success': False,
                    'error': "No PDF documents found or processed",
                    'documents_processed': 0,
                    'chunks_created': 0
                }
            
            # Add chunks to vector store
            logger.info(f"Adding {len(chunks)} chunks to vector store...")
            success = self.vector_store.add_documents(chunks)
            
            if success:
                # Get updated collection info
                updated_info = self.vector_store.get_collection_info()
                
                logger.info(f"Successfully loaded {len(chunks)} chunks from PDFs")
                return {
                    'success': True,
                    'message': f"Successfully loaded {len(chunks)} chunks",
                    'documents_processed': len(set(chunk['source_file'] for chunk in chunks)),
                    'chunks_created': len(chunks),
                    'collection_info': updated_info
                }
            else:
                return {
                    'success': False,
                    'error': "Failed to add documents to vector store",
                    'documents_processed': len(set(chunk['source_file'] for chunk in chunks)),
                    'chunks_created': len(chunks)
                }
                
        except Exception as e:
            logger.error(f"Error loading documents: {str(e)}")
            return {
                'success': False,
                'error': str(e),
                'documents_processed': 0,
                'chunks_created': 0
            }
    
    def get_processing_stats(self) -> Dict:
        """Get statistics about the document processing."""
        try:
            collection_info = self.vector_store.get_collection_info()
            
            # Count unique source files
            if collection_info.get('document_count', 0) > 0:
                # This would require querying the collection to get unique sources
                # For now, return basic info
                return {
                    'total_chunks': collection_info.get('document_count', 0),
                    'collection_name': collection_info.get('collection_name', ''),
                    'vector_db_path': collection_info.get('vector_db_path', ''),
                    'pdf_directory': self.config.PDF_DIRECTORY,
                    'pdf_directory_exists': os.path.exists(self.config.PDF_DIRECTORY)
                }
            else:
                return {
                    'total_chunks': 0,
                    'collection_name': collection_info.get('collection_name', ''),
                    'vector_db_path': collection_info.get('vector_db_path', ''),
                    'pdf_directory': self.config.PDF_DIRECTORY,
                    'pdf_directory_exists': os.path.exists(self.config.PDF_DIRECTORY)
                }
                
        except Exception as e:
            logger.error(f"Error getting processing stats: {str(e)}")
            return {
                'error': str(e),
                'total_chunks': 0
            }
    
    def validate_pdf_directory(self) -> Dict:
        """Validate the PDF directory and its contents."""
        try:
            if not os.path.exists(self.config.PDF_DIRECTORY):
                return {
                    'valid': False,
                    'error': f"PDF directory {self.config.PDF_DIRECTORY} does not exist",
                    'pdf_files': [],
                    'total_files': 0
                }
            
            pdf_files = [f for f in os.listdir(self.config.PDF_DIRECTORY) if f.lower().endswith('.pdf')]
            
            return {
                'valid': True,
                'pdf_directory': self.config.PDF_DIRECTORY,
                'pdf_files': pdf_files,
                'total_files': len(pdf_files)
            }
            
        except Exception as e:
            logger.error(f"Error validating PDF directory: {str(e)}")
            return {
                'valid': False,
                'error': str(e),
                'pdf_files': [],
                'total_files': 0
            }