src/storage/vector_database.py

"""
Vector Database Management for Legal RAG Pipeline
================================================

This module provides comprehensive vector database functionality for the Legal
RAG Pipeline, including document loading, text extraction, OCR processing,
and FAISS-based vector storage with GPU acceleration support.

Features:
- Multi-format document loading (PDF, web pages, text files)
- OCR fallback for scanned documents and images
- FAISS vector store with GPU acceleration
- Chunking and embedding of documents
- Web document scraping and processing
- Local file caching and management


"""

from langchain_community.document_loaders import WebBaseLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import faiss
try:
    from faiss import StandardGpuResources, index_cpu_to_gpu
except ImportError:
    StandardGpuResources = None
    index_cpu_to_gpu = None
from langchain_community.vectorstores import FAISS
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from src.generation.model_provider import get_embedding_model
from langchain_core.documents import Document
import os
import pytesseract
from pdf2image import convert_from_path
import webbrowser
from PIL import Image
import pytesseract
import requests
import time
from src.generation.model_provider import get_embedding_model
import numpy as np
from uuid import uuid4
from pathlib import Path
from tqdm import tqdm
import logging

# Set up logger
logger = logging.getLogger(__name__)

# Ensure local directory exists for storing downloaded PDFs
local_file_dir = "app_data/local_pdfs"
os.makedirs(local_file_dir, exist_ok=True)


def extract_text_with_ocr(file_path):
    """
    Extract text from PDF or image files using OCR technology.
    
    This function uses Tesseract OCR to extract text from scanned documents
    and images when standard text extraction methods fail.
    
    Args:
        file_path (str): Path to the PDF or image file
        
    Returns:
        str: Extracted text content
    """
    ext = os.path.splitext(file_path)[1].lower()
    text = ""
    
    try:
        if ext == ".pdf":
            # Convert PDF pages to images and apply OCR
            images = convert_from_path(file_path)
            for img in images:
                text += pytesseract.image_to_string(img)
        elif ext in [".png", ".jpg", ".jpeg", ".bmp", ".tiff"]:
            # Apply OCR directly to image files
            img = Image.open(file_path)
            text = pytesseract.image_to_string(img)
    except Exception as e:
        logger.info(f"OCR extraction failed for {file_path}: {e}")
    
    return text


def load_source_with_fallback(source):
    """
    Load document text with automatic fallback to OCR if needed.
    
    This function attempts to load documents using standard methods
    (PyPDFLoader for PDFs, web scraping for URLs) and falls back to
    OCR processing if standard extraction fails or produces poor results.
    
    Supports multiple document formats:
    - PDF files (local and web URLs)
    - Image files (PNG, JPG, JPEG, BMP, TIFF)
    - Text files (UTF-8 encoded)
    
    Args:
        source (str): Path to the document file or URL
        
    Returns:
        Document: LangChain Document object containing extracted text and metadata
    """
    doc = None
    local_file_path = os.path.join(local_file_dir, source.split("/")[-1])
    ext = os.path.splitext(source)[1].lower()
    previous_cache_from_web = f"{local_file_path}.txt"

    if os.path.exists(previous_cache_from_web):
        source = previous_cache_from_web
        ext = ".txt"

    if ext == ".pdf":
        # Handle web PDF downloads
        if source.startswith("https"):
            logger.info(f"\tLoading PDF from URL: {source}")
            # Download the PDF if it doesn't exist locally
            if not os.path.exists(local_file_path):
                headers = {
                            "User-Agent": "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:139.0) Gecko/20100101 Firefox/139.0",
                            }

                response = requests.get(source, headers=headers, stream=True)
                if response.status_code == 200:
                    os.makedirs(local_file_dir, exist_ok=True)
                    with open(local_file_path, "wb") as pdf_file:
                        for chunk in response.iter_content(chunk_size=1024):
                            pdf_file.write(chunk)
                    logger.info(f"PDF downloaded successfully to {local_file_path}")
                    ocr_text = extract_text_with_ocr(local_file_path)
                    if ocr_text.strip():
                        doc = Document(page_content=ocr_text, metadata={'source': source, "ocr_applied": True})
    
                elif response.status_code == 404:
                    source = source.replace(".pdf", ".html")
                    web_doc = WebBaseLoader(source, continue_on_failure=True).load()[0]
                    web_doc.page_content = web_doc.page_content.strip()
                    with open(previous_cache_from_web, "w", encoding="utf-8") as f:
                        f.write(web_doc.page_content)
                    return web_doc

                elif response.status_code == 429:
                    time.sleep(10)  # Wait for the browser to close
                    return load_source_with_fallback(source)

            else:
                logger.info(f"PDF already exists locally at {local_file_path}")
                ocr_text = extract_text_with_ocr(local_file_path)
                if ocr_text.strip():
                    doc = Document(page_content=ocr_text, metadata={'source': source, "ocr_applied": True})
    
        else:
            logger.info(f"\tUsing PyPDFLoader for local PDF")
            loader = PyPDFLoader(source)
            doc = loader.load()[0]

        # If no text extracted, fallback to OCR
        if not doc or not doc.page_content.strip():
            logger.info(f"\tNo text extracted from PDF, applying OCR")
            ocr_text = extract_text_with_ocr(source)
            if ocr_text.strip():
                doc = Document(page_content=ocr_text, metadata={'source': source, "ocr_applied": True})
    elif ext in [".png", ".jpg", ".jpeg", ".bmp", ".tiff"]:
        ocr_text = extract_text_with_ocr(source)
        if ocr_text.strip():
            doc = Document(page_content=ocr_text, metadata={'source': source, "ocr_applied": True})
    else:
        # Fallback: treat as text file
        logger.info(f"\tLoading text file: {source}")
        with open(source, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        doc = Document(page_content=content, metadata={'source': source})
    return doc

# Index documents from URLs and create a retriever for RAG

class VectorStoreManager:
    """
    Vector Store Manager for Legal Document Processing.
    
    This class manages FAISS vector stores with GPU acceleration support,
    handling document indexing, embedding generation, and similarity search
    for legal document retrieval.
    
    Features:
    - GPU-accelerated FAISS indexing with configurable memory allocation
    - Multiple document format support (PDF, web, text)
    - Automatic embedding dimension detection
    - IVF (Inverted File) indexing for large-scale similarity search
    - Persistent storage and loading of vector indexes
    - Memory optimization for resource-constrained environments
    
    Attributes:
        file_path (str): Path to vector store file for persistence
        use_gpu (bool): Whether to use GPU acceleration
        nlist (int): Number of clusters for IVF index
        gpu_memory_mb (int): GPU memory allocation in MB
        gpu_res: FAISS GPU resources object
        embeddings: Embedding model instance
        dim (int): Embedding vector dimensionality
        vector_store: FAISS vector store instance
    """
    
    def __init__(self, file_path=None, nlist=100, gpu_memory_mb=128):
        """
        Initialize VectorStoreManager with configurable GPU and memory settings.
        
        Sets up embedding model, detects vector dimensions, and initializes
        GPU resources.
        
        Args:
            file_path (str, optional): Path to existing vector store file
            nlist (int): Number of clusters for IVF index
            gpu_memory_mb (int): GPU memory allocation in MB
        """
        self.file_path = file_path
        self.nlist = nlist
        self.gpu_memory_mb = gpu_memory_mb
        self.gpu_res = None
        self.batch_size = 128
        
        # Configure embedding model with optimized settings
        self.embeddings_config = {"wait_time": 60}
        self.embeddings = get_embedding_model("mistral-embed", **self.embeddings_config)

        # Detect embedding dimensionality
        test_embedding = self.embeddings.embed_query("Test query to determine embedding dimensionality.")
        self.dim = len(test_embedding)
        logger.info(f"Embedding dimension detected: {self.dim}")


        # Initialize GPU resources when available. Hugging Face Spaces and
        # most local installs use faiss-cpu, which does not expose GPU helpers.
        if StandardGpuResources is not None:
            try:
                self.gpu_res = StandardGpuResources()
                self.gpu_res.setTempMemory(self.gpu_memory_mb * 1024 * 1024)
                logger.info(f"GPU resources initialized for FAISS with {self.gpu_memory_mb} MB memory allocation.")
            except Exception as e:
                logger.info(f"FAISS GPU initialization failed; using CPU index: {e}")
                self.gpu_res = None
        else:
            logger.info("FAISS GPU helpers are unavailable; using CPU index.")
          
        if self.file_path: # Load existing vector store from disk
            self.vector_store = FAISS.load_local(
                self.file_path, self.embeddings, allow_dangerous_deserialization=True
            )
            self._move_to_gpu()
        else: # Create a new vector store
            # Start with 'Flat' and upgrade to IVF when dataset size exceeds threshold
            index = self._create_optimized_index()
            self.vector_store = FAISS(
                embedding_function=self.embeddings,
                index=index,
                docstore=InMemoryDocstore(),
                index_to_docstore_id={},
            )

        self.text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
            chunk_size=1500,
            chunk_overlap=150, # 10% overlap
            separators=["\n\n", "\n", ".", "!", "?", "。", "！", "？", " "]
        )

    def _train_index(self, embedding_vectors: list[np.ndarray]):
        """
        Train the FAISS index if it requires training (IVF indices).
        Also handles switching from Flat to IVF index when enough vectors are available.
        
        Args:
            embedding_vectors: List of embedding vectors for training
        """
        index = self.vector_store.index
        num_vectors = len(embedding_vectors)

        # Check if we have enough vectors to qualify for IVF
        ivf_index_required = num_vectors + index.ntotal >= self.nlist * 2 
        
        # Upgrade
        if ("Flat" in type(index).__name__ and 
            ivf_index_required):  # Check if we have enough vectors to qualify for IVF
            logger.info(f"Switching to IVF index with {num_vectors + index.ntotal} vectors (>= {self.nlist * 2})")
            self._upgrade_to_ivf_index()
            index = self.vector_store.index
        
            # Check if we need to train the index
            if hasattr(index, "is_trained") and not index.is_trained:
                # Check if we have enough vectors to train IVF index
                if ivf_index_required:
                    logger.info(f"Not enough vectors to train IVF index with {self.nlist} clusters. "
                        f"Need at least {self.nlist} vectors. Deferring training.")
                    return
                
            vectors_np = np.array(embedding_vectors, dtype='float32')
            logger.info(f"Training index with {len(vectors_np)} vectors...")
            
            # Training happens on the current device (GPU or CPU)
            index.train(vectors_np)
            logger.info("Index training completed.")
     
    def _move_to_gpu(self):
        """
        Move the loaded vector store index to GPU.
        Handles different index types that might be loaded from disk.
        """
        if self.gpu_res:
            current_index = self.vector_store.index
            
            # Check if index is already on GPU
            if "Gpu" in type(current_index).__name__:
                logger.info("Index is already on GPU.")
                return
                
            try:
                # Move index to GPU
                gpu_index = index_cpu_to_gpu(self.gpu_res, 0, current_index)
                self.vector_store.index = gpu_index
                logger.info(f"Vector store index moved to GPU after load: {type(gpu_index).__name__}")
            except Exception as e:
                logger.info(f"Failed to move index to GPU (will continue with CPU): {e}")

    def _upgrade_to_ivf_index(self):
        """
        Upgrade from a Flat index to an IVF index when we have enough vectors.
        """
        current_index = self.vector_store.index
        
        # Extract all vectors from current index
        if current_index.ntotal > 0:
            logger.info(f"Extracting {current_index.ntotal} vectors from current index...")
            vectors = np.zeros((current_index.ntotal, self.dim), dtype='float32')
            for i in range(current_index.ntotal):
                vectors[i] = current_index.reconstruct(i)
        else:
            vectors = np.empty((0, self.dim), dtype='float32')
        
        # Create new IVF index
        quantizer = faiss.IndexFlatL2(self.dim)
        new_index = faiss.IndexIVFFlat(quantizer, self.dim, self.nlist)
        
        # Move to GPU if needed
        if self.gpu_res:
            try:
                new_index = index_cpu_to_gpu(self.gpu_res, 0, new_index)
                logger.info("New IVF index moved to GPU.")
            except Exception as e:
                logger.info(f"Failed to move new IVF index to GPU: {e}")
                logger.info("Continuing with CPU-based IVF index.")
                # Don't disable GPU entirely, just use CPU for this index
        
        # Train the new index if we have enough vectors
        if len(vectors) >= self.nlist:
            new_index.train(vectors)
            logger.info(f"New IVF index trained with {len(vectors)} vectors.")
        
        # Add existing vectors to new index
        if len(vectors) > 0:
            new_index.add(vectors)
            logger.info(f"Added {len(vectors)} existing vectors to new IVF index.")
        
        # Replace the index in vector store
        self.vector_store.index = new_index
        logger.info("Successfully upgraded to IVF index.")

    def _index_document(self, content, doc_metadata={'type': "unknown"}):
        """
        Index a single document by splitting it into chunks and adding to vector store.
        
        Args:
            content: Document content to index
            doc_metadata: Metadata for the document
            
        Returns:
            List of chunk UUIDs
        """
        doc_chunks = self.text_splitter.split_text(content)
        chunk_uuids = [str(uuid4()) for _ in range(len(doc_chunks))]
        
        # Clean up chunks to remove any empty strings
        doc_chunks = [chunk for chunk in doc_chunks if chunk.strip()]

        
        # Generate embeddings for all chunks
        vectors = [self.embeddings.embed_query(chunk) for chunk in doc_chunks]

        try:
            # Train the index if necessary (only for IVF indices)
            self._train_index(vectors)
            
            # Add texts with appropriate batch size for GPU performance
            self.vector_store.add_texts(
                texts=doc_chunks,
                metadatas=[doc_metadata] * len(doc_chunks),
                ids=chunk_uuids,
                batch_size=self.batch_size
            )
            logger.info(f"Successfully indexed {len(doc_chunks)} chunks from {doc_metadata.get('type', 'unknown')} document.")
        except Exception as e:
            logger.info(f"Error adding document {doc_metadata.get('type', 'unknown')} to vector store: {e}")
        return chunk_uuids

    def save_vector_store(self, index_path=None):
        if not index_path:
            if not self.file_path:
                timestamp = int(time.time())
                index_path = os.path.join("data/vector_db", f"vector_store_{timestamp}")
                os.makedirs(index_path, exist_ok=True)
            else:
                index_path = self.file_path
        else:
            os.makedirs(index_path, exist_ok=True)

        # Move index to CPU before saving if on GPU
        index_on_gpu = False
        current_index = self.vector_store.index
        
        # Check if index is on GPU by examining its type name
        if "Gpu" in type(current_index).__name__:
            try:
                import faiss
                cpu_index = faiss.index_gpu_to_cpu(current_index)
                self.vector_store.index = cpu_index
                index_on_gpu = True
                logger.info("Moved index to CPU for saving.")
            except Exception as e:
                logger.info(f"Failed to move index to CPU for saving: {e}")

        self.vector_store.save_local(index_path)
        self.file_path = index_path
        logger.info(f"Saved vector store to: {index_path}")

        # Move back to GPU if needed
        if index_on_gpu and self.gpu_res:
            try:
                from faiss import index_cpu_to_gpu
                self.vector_store.index = index_cpu_to_gpu(self.gpu_res, 0, self.vector_store.index)
                logger.info("Restored index to GPU after saving.")
            except Exception as e:
                logger.info(f"Failed to restore index to GPU after saving: {e}")
        return index_path

    def index_file_documents(self, paths, source_type):
        meta_datas = []
        for path in tqdm(paths, desc="Indexing file documents"):
            doc = load_source_with_fallback(path)
            if not doc:
                logger.info(f"Failed to load document from {path}. Skipping.")
                continue
            doc.metadata["type"] = source_type
            doc.metadata["chunk_ids"] = self._index_document(doc.page_content, doc.metadata)
            meta_datas.append(doc.metadata)
        return meta_datas
    
    def index_web_documents(self, urls, source_type):
        web_documents = WebBaseLoader(urls, continue_on_failure=True).load()
        # Process each URL
        meta_datas = []
        for doc in web_documents:
            # Doc metadata set by WebBaseLoader: 
            #   "source": url
            #   "title": soup.find("title").get_text()
            #   "description": soup.find("meta", attrs={"name": "description"}).get("content", "No description found.")
            #   "language": soup.find("html").get("lang", "No language found.")
            doc.metadata["source_type"] = source_type
            doc.metadata["chunk_ids"] = self._index_document(doc.page_content, doc.metadata)
            meta_datas.append(doc.metadata)
        return meta_datas

    def add_texts(self, texts, metadatas=None, ids=None):
        """
        Add texts to the vector store with proper training if needed.
        Optimized for GPU performance when available.
        
        Args:
            texts: List of text strings to add
            metadatas: List of metadata dictionaries (optional)
            ids: List of document IDs (optional)
            
        Returns:
            List of document IDs
        """
        if not ids:
            ids = [str(uuid4()) for _ in range(len(texts))]
        
        # Generate embeddings for training if needed
        logger.info(f"Generating embeddings for {len(texts)} texts...")
        vectors = [self.embeddings.embed_query(text) for text in texts]
        
        # Train the index if necessary
        self._train_index(vectors)
        
        # Add the texts to the vector store
        result = self.vector_store.add_texts(
            texts=texts, metadatas=metadatas, ids=ids, batch_size=self.batch_size
        )
        
        logger.info(f"Successfully added {len(texts)} texts to vector store.")
        return result

    def query_vector_store(self, query, k=5, meta_filter=None):
        if not self.vector_store:
            logger.info("No vector store loaded.")
            return []
        return self.vector_store.similarity_search_with_score(query, k=k, filter=meta_filter)

    def _create_optimized_index(self):
        """
        Create an optimized FAISS index based on expected dataset size and GPU availability.
        Starts with a Flat index and upgrades to IVF when enough vectors are available.
        
        Returns:
            FAISS index optimized for the current configuration
        """
        # Always start with a Flat index for small datasets
        # We'll upgrade to IVF when we have enough vectors (>= nlist * 2)
        index = faiss.IndexFlatL2(self.dim)
        
        if self.gpu_res:
            try:
                # Move to GPU immediately for better performance
                index = index_cpu_to_gpu(self.gpu_res, 0, index)
                logger.info(f"Created Flat index on GPU (will upgrade to IVF when >= {self.nlist * 2} vectors)")
                return index
            except Exception as e:
                logger.info(f"Failed to create GPU index: {e}")
                logger.info("This might be due to GPU memory limitations. Falling back to CPU.")
                self.gpu_res = None
                
        # CPU index
        logger.info(f"Using CPU-based IndexFlatL2 (will upgrade to IVF when >= {self.nlist * 2} vectors)")
        return index

    def get_index_info(self):
        """
        Get information about the current FAISS index configuration.
        
        Returns:
            Dictionary with index information
        """
        if not self.vector_store or not self.vector_store.index:
            return {"status": "No index available"}
            
        index = self.vector_store.index
        info = {
            "index_type": type(index).__name__,
            "dimension": self.dim,
            "total_vectors": index.ntotal,
            "is_trained": getattr(index, "is_trained", True),
            "on_gpu": "Gpu" in type(index).__name__
        }
        
        if hasattr(index, "nlist"):
            info["nlist"] = index.nlist
            
        return info
    
    def print_index_status(self):
        """
        Print current index status and configuration.
        """
        info = self.get_index_info()
        logger.info("=== FAISS Index Status ===")
        for key, value in info.items():
            logger.info(f"{key}: {value}")
        logger.info("=" * 26)


    def get_case_vector_store_dict(self):
        """
        Get a dictionary representation of the vector store for case management.
        """
        return self.vector_store.docstore._dict

    def delete_documents(self, uuids_to_delete):
        """
        Delete documents from the vector store by their UUIDs.
        
        Args:
            uuids_to_delete (list): List of UUIDs to delete from the vector store
        """
        if not uuids_to_delete:
            logger.info("No UUIDs provided for deletion.")
            return
        
        logger.info(f"Deleting {len(uuids_to_delete)} documents from vector store...")
        self.vector_store.delete(uuids_to_delete)
        logger.info("Deletion completed.")