"""
Library RAG Service

Handles indexing and searching library documents using RAG:
- Index text documents into vector database
- Chunk documents for semantic search
- Generate embeddings using local models
- Manage FAISS indices per research
- Track RAG status in library
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from langchain_core.documents import Document as LangchainDocument
from loguru import logger
from sqlalchemy import func

from ...config.paths import get_cache_directory
from ...database.models.library import (
    Document,
    DocumentChunk,
    DocumentCollection,
    Collection,
    RAGIndex,
    RagDocumentStatus,
    EmbeddingProvider,
)
from ...database.session_context import get_user_db_session
from ...utilities.type_utils import to_bool
from ..utils import ensure_in_collection
from ...embeddings.splitters import get_text_splitter
from ...web_search_engines.engines.local_embedding_manager import (
    LocalEmbeddingManager,
)
from ...security.file_integrity import FileIntegrityManager, FAISSIndexVerifier
import hashlib
from faiss import IndexFlatL2, IndexFlatIP, IndexHNSWFlat
from langchain_community.vectorstores import FAISS
from langchain_community.docstore.in_memory import InMemoryDocstore


class LibraryRAGService:
    """Service for managing RAG indexing of library documents."""

    def __init__(
        self,
        username: str,
        embedding_model: str = "all-MiniLM-L6-v2",
        embedding_provider: str = "sentence_transformers",
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        splitter_type: str = "recursive",
        text_separators: Optional[list] = None,
        distance_metric: str = "cosine",
        normalize_vectors: bool = True,
        index_type: str = "flat",
        embedding_manager: Optional["LocalEmbeddingManager"] = None,
        db_password: Optional[str] = None,
    ):
        """
        Initialize library RAG service for a user.

        Args:
            username: Username for database access
            embedding_model: Name of the embedding model to use
            embedding_provider: Provider type ('sentence_transformers' or 'ollama')
            chunk_size: Size of text chunks for splitting
            chunk_overlap: Overlap between consecutive chunks
            splitter_type: Type of splitter ('recursive', 'token', 'sentence', 'semantic')
            text_separators: List of text separators for chunking (default: ["\n\n", "\n", ". ", " ", ""])
            distance_metric: Distance metric ('cosine', 'l2', or 'dot_product')
            normalize_vectors: Whether to normalize vectors with L2
            index_type: FAISS index type ('flat', 'hnsw', or 'ivf')
            embedding_manager: Optional pre-constructed LocalEmbeddingManager for testing/flexibility
            db_password: Optional database password for background thread access
        """
        self.username = username
        self._db_password = db_password  # Can be used for thread access
        # Initialize optional attributes to None before they're set below
        # This allows the db_password setter to check them without hasattr
        self.embedding_manager = None
        self.integrity_manager = None
        self.embedding_model = embedding_model
        self.embedding_provider = embedding_provider
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.splitter_type = splitter_type
        self.text_separators = (
            text_separators
            if text_separators is not None
            else ["\n\n", "\n", ". ", " ", ""]
        )
        self.distance_metric = distance_metric
        # Ensure normalize_vectors is always a proper boolean
        self.normalize_vectors = to_bool(normalize_vectors, default=True)
        self.index_type = index_type

        # Emit the active configuration so users can confirm their
        # UI-configured embedding settings are being honored (regression
        # signal for #3453).
        logger.info(
            f"RAG service initialized for user={username}: "
            f"provider={embedding_provider} model={embedding_model} "
            f"chunk_size={chunk_size} chunk_overlap={chunk_overlap} "
            f"splitter={splitter_type} index_type={index_type}"
        )

        # Use provided embedding manager or create a new one
        # (Must be created before text splitter for semantic chunking)
        if embedding_manager is not None:
            self.embedding_manager = embedding_manager
        else:
            # Initialize embedding manager with library collection
            # Load the complete user settings snapshot from database using the proper method
            from ...settings.manager import SettingsManager

            # Use proper database session for SettingsManager
            # Note: using _db_password (backing field) directly here because the
            # db_password property setter propagates to embedding_manager/integrity_manager,
            # which are still None at this point in __init__.
            with get_user_db_session(username, self._db_password) as session:
                settings_manager = SettingsManager(session)
                settings_snapshot = settings_manager.get_settings_snapshot()

            # Add the specific settings needed for this RAG service
            settings_snapshot.update(
                {
                    "_username": username,
                    "embeddings.provider": embedding_provider,
                    f"embeddings.{embedding_provider}.model": embedding_model,
                    "local_search_chunk_size": chunk_size,
                    "local_search_chunk_overlap": chunk_overlap,
                }
            )

            self.embedding_manager = LocalEmbeddingManager(
                embedding_model=embedding_model,
                embedding_model_type=embedding_provider,
                settings_snapshot=settings_snapshot,
            )

        # Initialize text splitter based on type
        # (Must be created AFTER embedding_manager for semantic chunking)
        self.text_splitter = get_text_splitter(
            splitter_type=self.splitter_type,
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            text_separators=self.text_separators,
            embeddings=self.embedding_manager.embeddings
            if self.splitter_type == "semantic"
            else None,
        )

        # Initialize or load FAISS index for library collection
        self.faiss_index = None
        self.rag_index_record = None

        # Initialize file integrity manager for FAISS indexes
        self.integrity_manager = FileIntegrityManager(
            username, password=self._db_password
        )
        self.integrity_manager.register_verifier(FAISSIndexVerifier())

        self._closed = False

    def close(self):
        """Release embedding model and index resources."""
        if self._closed:
            return
        self._closed = True

        # Clear embedding manager resources
        if self.embedding_manager is not None:
            # Clear references to allow garbage collection
            self.embedding_manager = None

        # Clear FAISS index
        if self.faiss_index is not None:
            self.faiss_index = None

        # Clear other resources
        self.rag_index_record = None
        self.integrity_manager = None
        self.text_splitter = None

    def __enter__(self):
        """Enter context manager."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager, ensuring cleanup."""
        self.close()
        return False

    @property
    def db_password(self):
        """Get database password."""
        return self._db_password

    @db_password.setter
    def db_password(self, value):
        """Set database password and propagate to embedding manager and integrity manager."""
        self._db_password = value
        if self.embedding_manager:
            self.embedding_manager.db_password = value
        if self.integrity_manager:
            self.integrity_manager.password = value

    def _get_index_hash(
        self,
        collection_name: str,
        embedding_model: str,
        embedding_model_type: str,
    ) -> str:
        """Generate hash for index identification."""
        hash_input = (
            f"{collection_name}:{embedding_model}:{embedding_model_type}"
        )
        return hashlib.sha256(hash_input.encode()).hexdigest()

    def _get_index_path(self, index_hash: str) -> Path:
        """Get path for FAISS index file."""
        # Store in centralized cache directory (respects LDR_DATA_DIR)
        cache_dir = get_cache_directory() / "rag_indices"
        cache_dir.mkdir(parents=True, exist_ok=True)
        return cache_dir / f"{index_hash}.faiss"

    @staticmethod
    def _deduplicate_chunks(
        chunks: List[LangchainDocument],
        chunk_ids: List[str],
        existing_ids: Optional[set] = None,
    ) -> Tuple[List[LangchainDocument], List[str]]:
        """Deduplicate chunks by ID within a batch, optionally excluding existing IDs."""
        seen_ids: set = set()
        new_chunks: List[LangchainDocument] = []
        new_ids: List[str] = []
        for chunk, chunk_id in zip(chunks, chunk_ids):
            if chunk_id not in seen_ids and (
                existing_ids is None or chunk_id not in existing_ids
            ):
                new_chunks.append(chunk)
                new_ids.append(chunk_id)
                seen_ids.add(chunk_id)
        return new_chunks, new_ids

    def _get_or_create_rag_index(self, collection_id: str) -> RAGIndex:
        """Get or create RAGIndex record for the current configuration."""
        with get_user_db_session(self.username, self.db_password) as session:
            # Use collection_<uuid> format
            collection_name = f"collection_{collection_id}"
            index_hash = self._get_index_hash(
                collection_name, self.embedding_model, self.embedding_provider
            )

            # Try to get existing index
            rag_index = (
                session.query(RAGIndex).filter_by(index_hash=index_hash).first()
            )

            if not rag_index:
                # Create new index record
                index_path = self._get_index_path(index_hash)

                # Get embedding dimension by embedding a test string
                test_embedding = self.embedding_manager.embeddings.embed_query(
                    "test"
                )
                embedding_dim = len(test_embedding)

                rag_index = RAGIndex(
                    collection_name=collection_name,
                    embedding_model=self.embedding_model,
                    embedding_model_type=EmbeddingProvider(
                        self.embedding_provider
                    ),
                    embedding_dimension=embedding_dim,
                    index_path=str(index_path),
                    index_hash=index_hash,
                    chunk_size=self.chunk_size,
                    chunk_overlap=self.chunk_overlap,
                    splitter_type=self.splitter_type,
                    text_separators=self.text_separators,
                    distance_metric=self.distance_metric,
                    normalize_vectors=self.normalize_vectors,
                    index_type=self.index_type,
                    chunk_count=0,
                    total_documents=0,
                    status="active",
                    is_current=True,
                )
                session.add(rag_index)
                session.commit()
                session.refresh(rag_index)
                logger.info(f"Created new RAG index: {index_hash}")

            return rag_index

    def load_or_create_faiss_index(self, collection_id: str) -> FAISS:
        """
        Load existing FAISS index or create new one.

        Args:
            collection_id: UUID of the collection

        Returns:
            FAISS vector store instance
        """
        rag_index = self._get_or_create_rag_index(collection_id)
        self.rag_index_record = rag_index

        index_path = Path(rag_index.index_path)

        if index_path.exists():
            # Verify integrity before loading
            verified, reason = self.integrity_manager.verify_file(index_path)
            if not verified:
                logger.error(
                    f"Integrity verification failed for {index_path}: {reason}. "
                    f"Refusing to load. Creating new index."
                )
                # Remove corrupted index
                try:
                    index_path.unlink()
                    logger.info(f"Removed corrupted index file: {index_path}")
                except Exception:
                    logger.exception("Failed to remove corrupted index")
            else:
                try:
                    # Check for embedding dimension mismatch before loading
                    current_dim = len(
                        self.embedding_manager.embeddings.embed_query(
                            "dimension_check"
                        )
                    )
                    stored_dim = rag_index.embedding_dimension

                    if stored_dim and current_dim != stored_dim:
                        logger.warning(
                            f"Embedding dimension mismatch detected! "
                            f"Index created with dim={stored_dim}, "
                            f"current model returns dim={current_dim}. "
                            f"Deleting old index and rebuilding."
                        )
                        # Delete old index files
                        try:
                            index_path.unlink()
                            pkl_path = index_path.with_suffix(".pkl")
                            if pkl_path.exists():
                                pkl_path.unlink()
                            logger.info(
                                f"Deleted old FAISS index files at {index_path}"
                            )
                        except Exception:
                            logger.exception("Failed to delete old index files")

                        # Update RAGIndex with new dimension and reset counts
                        with get_user_db_session(
                            self.username, self.db_password
                        ) as session:
                            idx = (
                                session.query(RAGIndex)
                                .filter_by(id=rag_index.id)
                                .first()
                            )
                            if idx:
                                idx.embedding_dimension = current_dim
                                idx.chunk_count = 0
                                idx.total_documents = 0
                                session.commit()
                                logger.info(
                                    f"Updated RAGIndex dimension to {current_dim}"
                                )

                            # Clear rag_document_status for this index
                            session.query(RagDocumentStatus).filter_by(
                                rag_index_id=rag_index.id
                            ).delete()
                            session.commit()
                            logger.info(
                                "Cleared indexed status for documents in this "
                                "collection"
                            )

                        # Update local reference for index creation below
                        rag_index.embedding_dimension = current_dim
                        # Fall through to create new index below
                    else:
                        # Dimensions match (or no stored dimension), load index
                        faiss_index = FAISS.load_local(
                            str(index_path.parent),
                            self.embedding_manager.embeddings,
                            index_name=index_path.stem,
                            allow_dangerous_deserialization=True,
                            normalize_L2=True,
                        )
                        logger.info(
                            f"Loaded existing FAISS index from {index_path}"
                        )
                        return faiss_index
                except Exception:
                    logger.warning(
                        "Failed to load FAISS index, creating new one"
                    )

        # Create new FAISS index with configurable type and distance metric
        logger.info(
            f"Creating new FAISS index: type={self.index_type}, metric={self.distance_metric}, dimension={rag_index.embedding_dimension}"
        )

        # Create index based on type and distance metric
        if self.index_type == "hnsw":
            # HNSW: Fast approximate search, best for large collections
            # M=32 is a good default for connections per layer
            index = IndexHNSWFlat(rag_index.embedding_dimension, 32)
            logger.info("Created HNSW index with M=32 connections")
        elif self.index_type == "ivf":
            # IVF requires training, for now fall back to flat
            # TODO: Implement IVF with proper training
            logger.warning(
                "IVF index type not yet fully implemented, using Flat index"
            )
            if self.distance_metric in ("cosine", "dot_product"):
                index = IndexFlatIP(rag_index.embedding_dimension)
            else:
                index = IndexFlatL2(rag_index.embedding_dimension)
        else:  # "flat" or default
            # Flat index: Exact search
            if self.distance_metric in ("cosine", "dot_product"):
                # For cosine similarity, use inner product (IP) with normalized vectors
                index = IndexFlatIP(rag_index.embedding_dimension)
                logger.info(
                    "Created Flat index with Inner Product (for cosine similarity)"
                )
            else:  # l2
                index = IndexFlatL2(rag_index.embedding_dimension)
                logger.info("Created Flat index with L2 distance")

        faiss_index = FAISS(
            self.embedding_manager.embeddings,
            index=index,
            docstore=InMemoryDocstore(),  # Minimal - chunks in DB
            index_to_docstore_id={},
            normalize_L2=self.normalize_vectors,  # Use configurable normalization
        )
        logger.info(
            f"FAISS index created with normalization={self.normalize_vectors}"
        )
        return faiss_index

    def get_current_index_info(
        self, collection_id: Optional[str] = None
    ) -> Optional[Dict[str, Any]]:
        """
        Get information about the current RAG index for a collection.

        Args:
            collection_id: UUID of collection (defaults to Library if None)
        """
        with get_user_db_session(self.username, self.db_password) as session:
            # Get collection name in the format stored in RAGIndex (collection_<uuid>)
            if collection_id:
                collection = (
                    session.query(Collection)
                    .filter_by(id=collection_id)
                    .first()
                )
                collection_name = (
                    f"collection_{collection_id}" if collection else "unknown"
                )
            else:
                # Default to Library collection
                from ...database.library_init import get_default_library_id

                collection_id = get_default_library_id(
                    self.username, self.db_password
                )
                collection_name = f"collection_{collection_id}"

            rag_index = (
                session.query(RAGIndex)
                .filter_by(collection_name=collection_name, is_current=True)
                .first()
            )

            if not rag_index:
                # Debug: check all RAG indices for this collection
                all_indices = session.query(RAGIndex).all()
                logger.info(
                    f"No RAG index found for collection_name='{collection_name}'. All indices: {[(idx.collection_name, idx.is_current) for idx in all_indices]}"
                )
                return None

            # Calculate actual counts from rag_document_status table
            from ...database.models.library import RagDocumentStatus

            actual_chunk_count = (
                session.query(func.sum(RagDocumentStatus.chunk_count))
                .filter_by(collection_id=collection_id)
                .scalar()
                or 0
            )

            actual_doc_count = (
                session.query(RagDocumentStatus)
                .filter_by(collection_id=collection_id)
                .count()
            )

            return {
                "embedding_model": rag_index.embedding_model,
                "embedding_model_type": rag_index.embedding_model_type.value
                if rag_index.embedding_model_type
                else None,
                "embedding_dimension": rag_index.embedding_dimension,
                "chunk_size": rag_index.chunk_size,
                "chunk_overlap": rag_index.chunk_overlap,
                "chunk_count": actual_chunk_count,
                "total_documents": actual_doc_count,
                "created_at": rag_index.created_at.isoformat(),
                "last_updated_at": rag_index.last_updated_at.isoformat(),
            }

    def index_document(
        self, document_id: str, collection_id: str, force_reindex: bool = False
    ) -> Dict[str, Any]:
        """
        Index a single document into RAG for a specific collection.

        Args:
            document_id: UUID of the Document to index
            collection_id: UUID of the Collection to index for
            force_reindex: Whether to force reindexing even if already indexed

        Returns:
            Dict with status, chunk_count, and any errors
        """
        with get_user_db_session(self.username, self.db_password) as session:
            # Get the document
            document = session.query(Document).filter_by(id=document_id).first()

            if not document:
                return {"status": "error", "error": "Document not found"}

            # Get or create DocumentCollection entry
            doc_collection = ensure_in_collection(
                session, document_id, collection_id
            )

            # Check if already indexed for this collection
            if doc_collection.indexed and not force_reindex:
                return {
                    "status": "skipped",
                    "message": "Document already indexed for this collection",
                    "chunk_count": doc_collection.chunk_count,
                }

            # Validate text content
            if not document.text_content:
                return {
                    "status": "error",
                    "error": "Document has no text content",
                }

            try:
                # Create LangChain Document from text
                doc = LangchainDocument(
                    page_content=document.text_content,
                    metadata={
                        "source": document.original_url,
                        "document_id": document_id,  # Add document ID for source linking
                        "collection_id": collection_id,  # Add collection ID
                        "title": document.title
                        or document.filename
                        or "Untitled",
                        "document_title": document.title
                        or document.filename
                        or "Untitled",  # Add for compatibility
                        "authors": document.authors,
                        "published_date": str(document.published_date)
                        if document.published_date
                        else None,
                        "doi": document.doi,
                        "arxiv_id": document.arxiv_id,
                        "pmid": document.pmid,
                        "pmcid": document.pmcid,
                        "extraction_method": document.extraction_method,
                        "word_count": document.word_count,
                    },
                )

                # Split into chunks
                chunks = self.text_splitter.split_documents([doc])
                logger.info(
                    f"Split document {document_id} into {len(chunks)} chunks"
                )

                # Get collection name for chunk storage
                collection = (
                    session.query(Collection)
                    .filter_by(id=collection_id)
                    .first()
                )
                # Use collection_<uuid> format for internal storage
                collection_name = (
                    f"collection_{collection_id}" if collection else "unknown"
                )

                # Store chunks in database using embedding manager
                embedding_ids = self.embedding_manager._store_chunks_to_db(
                    chunks=chunks,
                    collection_name=collection_name,
                    source_type="document",
                    source_id=document_id,
                )

                # Load or create FAISS index
                if self.faiss_index is None:
                    self.faiss_index = self.load_or_create_faiss_index(
                        collection_id
                    )

                # If force_reindex, remove old chunks from FAISS before adding new ones
                if force_reindex:
                    existing_ids = (
                        set(self.faiss_index.docstore._dict.keys())
                        if hasattr(self.faiss_index, "docstore")
                        else set()
                    )
                    old_chunk_ids = list(
                        {eid for eid in embedding_ids if eid in existing_ids}
                    )
                    if old_chunk_ids:
                        logger.info(
                            f"Force re-index: removing {len(old_chunk_ids)} existing chunks from FAISS"
                        )
                        self.faiss_index.delete(old_chunk_ids)

                # Filter out chunks that already exist in FAISS (unless force_reindex)
                if not force_reindex:
                    existing_ids = (
                        set(self.faiss_index.docstore._dict.keys())
                        if hasattr(self.faiss_index, "docstore")
                        else set()
                    )
                else:
                    existing_ids = None

                unique_count = len(set(embedding_ids))
                batch_dups = len(chunks) - unique_count

                new_chunks, new_ids = self._deduplicate_chunks(
                    chunks, embedding_ids, existing_ids
                )

                # Add embeddings to FAISS index
                if new_chunks:
                    if force_reindex:
                        logger.info(
                            f"Force re-index: adding {len(new_chunks)} chunks with updated metadata to FAISS index"
                        )
                    else:
                        already_exist = unique_count - len(new_chunks)
                        logger.info(
                            f"Adding {len(new_chunks)} new embeddings to FAISS index "
                            f"({already_exist} already exist, {batch_dups} batch duplicates removed)"
                        )
                    self.faiss_index.add_documents(new_chunks, ids=new_ids)
                else:
                    logger.info(
                        f"All {len(chunks)} chunks already exist in FAISS index, skipping"
                    )

                # Save FAISS index
                index_path = Path(self.rag_index_record.index_path)
                self.faiss_index.save_local(
                    str(index_path.parent), index_name=index_path.stem
                )
                # Record file integrity
                self.integrity_manager.record_file(
                    index_path,
                    related_entity_type="rag_index",
                    related_entity_id=self.rag_index_record.id,
                )
                logger.info(
                    f"Saved FAISS index to {index_path} with integrity tracking"
                )

                from datetime import datetime, UTC
                from sqlalchemy import text

                # Check if document was already indexed (for stats update)
                existing_status = (
                    session.query(RagDocumentStatus)
                    .filter_by(
                        document_id=document_id, collection_id=collection_id
                    )
                    .first()
                )
                was_already_indexed = existing_status is not None

                # Mark document as indexed using rag_document_status table
                # Row existence = indexed, simple and clean
                timestamp = datetime.now(UTC)

                # Create or update RagDocumentStatus using ORM merge (atomic upsert)
                rag_status = RagDocumentStatus(
                    document_id=document_id,
                    collection_id=collection_id,
                    rag_index_id=self.rag_index_record.id,
                    chunk_count=len(chunks),
                    indexed_at=timestamp,
                )
                session.merge(rag_status)

                logger.info(
                    f"Marked document as indexed in rag_document_status: doc_id={document_id}, coll_id={collection_id}, chunks={len(chunks)}"
                )

                # Also update DocumentCollection table for backward compatibility
                session.query(DocumentCollection).filter_by(
                    document_id=document_id, collection_id=collection_id
                ).update(
                    {
                        "indexed": True,
                        "chunk_count": len(chunks),
                        "last_indexed_at": timestamp,
                    }
                )

                logger.info(
                    "Also updated DocumentCollection.indexed for backward compatibility"
                )

                # Update RAGIndex statistics (only if not already indexed)
                rag_index_obj = (
                    session.query(RAGIndex)
                    .filter_by(id=self.rag_index_record.id)
                    .first()
                )
                if rag_index_obj and not was_already_indexed:
                    rag_index_obj.chunk_count += len(chunks)
                    rag_index_obj.total_documents += 1
                    rag_index_obj.last_updated_at = datetime.now(UTC)
                    logger.info(
                        f"Updated RAGIndex stats: chunk_count +{len(chunks)}, total_documents +1"
                    )

                # Flush ORM changes to database before commit
                session.flush()
                logger.info(f"Flushed ORM changes for document {document_id}")

                # Commit the transaction
                session.commit()

                # WAL checkpoint after commit to ensure persistence
                session.execute(text("PRAGMA wal_checkpoint(FULL)"))

                logger.info(
                    f"Successfully indexed document {document_id} for collection {collection_id} "
                    f"with {len(chunks)} chunks"
                )

                return {
                    "status": "success",
                    "chunk_count": len(chunks),
                    "embedding_ids": embedding_ids,
                }

            except Exception as e:
                logger.exception(
                    f"Error indexing document {document_id} for collection {collection_id}"
                )
                return {
                    "status": "error",
                    "error": f"Operation failed: {type(e).__name__}",
                }

    def index_all_documents(
        self,
        collection_id: str,
        force_reindex: bool = False,
        progress_callback=None,
    ) -> Dict[str, Any]:
        """
        Index all documents in a collection into RAG.

        Args:
            collection_id: UUID of the collection to index
            force_reindex: Whether to force reindexing already indexed documents
            progress_callback: Optional callback function called after each document with (current, total, doc_title, status)

        Returns:
            Dict with counts of successful, skipped, and failed documents
        """
        with get_user_db_session(self.username, self.db_password) as session:
            # Get all DocumentCollection entries for this collection
            query = session.query(DocumentCollection).filter_by(
                collection_id=collection_id
            )

            if not force_reindex:
                # Only index documents that haven't been indexed yet
                query = query.filter_by(indexed=False)

            doc_collections = query.all()

            if not doc_collections:
                return {
                    "status": "info",
                    "message": "No documents to index",
                    "successful": 0,
                    "skipped": 0,
                    "failed": 0,
                }

            results = {"successful": 0, "skipped": 0, "failed": 0, "errors": []}
            total = len(doc_collections)

            for idx, doc_collection in enumerate(doc_collections, 1):
                # Get the document for title info
                document = (
                    session.query(Document)
                    .filter_by(id=doc_collection.document_id)
                    .first()
                )
                title = document.title if document else "Unknown"

                result = self.index_document(
                    doc_collection.document_id, collection_id, force_reindex
                )

                if result["status"] == "success":
                    results["successful"] += 1
                elif result["status"] == "skipped":
                    results["skipped"] += 1
                else:
                    results["failed"] += 1
                    results["errors"].append(
                        {
                            "doc_id": doc_collection.document_id,
                            "title": title,
                            "error": result.get("error"),
                        }
                    )

                # Call progress callback if provided
                if progress_callback:
                    progress_callback(idx, total, title, result["status"])

            logger.info(
                f"Indexed collection {collection_id}: "
                f"{results['successful']} successful, "
                f"{results['skipped']} skipped, "
                f"{results['failed']} failed"
            )

            return results

    def remove_document_from_rag(
        self, document_id: str, collection_id: str
    ) -> Dict[str, Any]:
        """
        Remove a document's chunks from RAG for a specific collection.

        Args:
            document_id: UUID of the Document to remove
            collection_id: UUID of the Collection to remove from

        Returns:
            Dict with status and count of removed chunks
        """
        with get_user_db_session(self.username, self.db_password) as session:
            # Get the DocumentCollection entry
            doc_collection = (
                session.query(DocumentCollection)
                .filter_by(document_id=document_id, collection_id=collection_id)
                .first()
            )

            if not doc_collection:
                return {
                    "status": "error",
                    "error": "Document not found in collection",
                }

            try:
                # Get collection name in the format collection_<uuid>
                collection = (
                    session.query(Collection)
                    .filter_by(id=collection_id)
                    .first()
                )
                # Use collection_<uuid> format for internal storage
                collection_name = (
                    f"collection_{collection_id}" if collection else "unknown"
                )

                # Delete chunks from database
                deleted_count = self.embedding_manager._delete_chunks_from_db(
                    collection_name=collection_name,
                    source_id=document_id,
                )

                # Update DocumentCollection RAG status
                doc_collection.indexed = False
                doc_collection.chunk_count = 0
                doc_collection.last_indexed_at = None
                session.commit()

                logger.info(
                    f"Removed {deleted_count} chunks for document {document_id} from collection {collection_id}"
                )

                return {"status": "success", "deleted_count": deleted_count}

            except Exception as e:
                logger.exception(
                    f"Error removing document {document_id} from collection {collection_id}"
                )
                return {
                    "status": "error",
                    "error": f"Operation failed: {type(e).__name__}",
                }

    def index_documents_batch(
        self,
        doc_info: List[tuple],
        collection_id: str,
        force_reindex: bool = False,
    ) -> Dict[str, Dict[str, Any]]:
        """
        Index multiple documents in a batch for a specific collection.

        Args:
            doc_info: List of (doc_id, title) tuples
            collection_id: UUID of the collection to index for
            force_reindex: Whether to force reindexing even if already indexed

        Returns:
            Dict mapping doc_id to individual result
        """
        results = {}
        doc_ids = [doc_id for doc_id, _ in doc_info]

        # Use single database session for querying
        with get_user_db_session(self.username, self.db_password) as session:
            # Pre-load all documents for this batch
            documents = (
                session.query(Document).filter(Document.id.in_(doc_ids)).all()
            )

            # Create lookup for quick access
            doc_lookup = {doc.id: doc for doc in documents}

            # Pre-load DocumentCollection entries
            doc_collections = (
                session.query(DocumentCollection)
                .filter(
                    DocumentCollection.document_id.in_(doc_ids),
                    DocumentCollection.collection_id == collection_id,
                )
                .all()
            )
            doc_collection_lookup = {
                dc.document_id: dc for dc in doc_collections
            }

            # Process each document in the batch
            for doc_id, title in doc_info:
                document = doc_lookup.get(doc_id)

                if not document:
                    results[doc_id] = {
                        "status": "error",
                        "error": "Document not found",
                    }
                    continue

                # Check if already indexed via DocumentCollection
                doc_collection = doc_collection_lookup.get(doc_id)
                if (
                    doc_collection
                    and doc_collection.indexed
                    and not force_reindex
                ):
                    results[doc_id] = {
                        "status": "skipped",
                        "message": "Document already indexed for this collection",
                        "chunk_count": doc_collection.chunk_count,
                    }
                    continue

                # Validate text content
                if not document.text_content:
                    results[doc_id] = {
                        "status": "error",
                        "error": "Document has no text content",
                    }
                    continue

                # Index the document
                try:
                    result = self.index_document(
                        doc_id, collection_id, force_reindex
                    )
                    results[doc_id] = result
                except Exception as e:
                    logger.exception(
                        f"Error indexing document {doc_id} in batch"
                    )
                    results[doc_id] = {
                        "status": "error",
                        "error": f"Indexing failed: {type(e).__name__}",
                    }

        return results

    def get_rag_stats(
        self, collection_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Get RAG statistics for a collection.

        Args:
            collection_id: UUID of the collection (defaults to Library)

        Returns:
            Dict with counts and metadata about indexed documents
        """
        with get_user_db_session(self.username, self.db_password) as session:
            # Get collection ID (default to Library)
            if not collection_id:
                from ...database.library_init import get_default_library_id

                collection_id = get_default_library_id(
                    self.username, self.db_password
                )

            # Count total documents in collection
            total_docs = (
                session.query(DocumentCollection)
                .filter_by(collection_id=collection_id)
                .count()
            )

            # Count indexed documents from rag_document_status table
            from ...database.models.library import RagDocumentStatus

            indexed_docs = (
                session.query(RagDocumentStatus)
                .filter_by(collection_id=collection_id)
                .count()
            )

            # Count total chunks from rag_document_status table
            total_chunks = (
                session.query(func.sum(RagDocumentStatus.chunk_count))
                .filter_by(collection_id=collection_id)
                .scalar()
                or 0
            )

            # Get collection name in the format stored in DocumentChunk (collection_<uuid>)
            collection = (
                session.query(Collection).filter_by(id=collection_id).first()
            )
            collection_name = (
                f"collection_{collection_id}" if collection else "library"
            )

            # Get embedding model info from chunks
            chunk_sample = (
                session.query(DocumentChunk)
                .filter_by(collection_name=collection_name)
                .first()
            )

            embedding_info = {}
            if chunk_sample:
                embedding_info = {
                    "model": chunk_sample.embedding_model,
                    "model_type": chunk_sample.embedding_model_type.value
                    if chunk_sample.embedding_model_type
                    else None,
                    "dimension": chunk_sample.embedding_dimension,
                }

            return {
                "total_documents": total_docs,
                "indexed_documents": indexed_docs,
                "unindexed_documents": total_docs - indexed_docs,
                "total_chunks": total_chunks,
                "embedding_info": embedding_info,
                "chunk_size": self.chunk_size,
                "chunk_overlap": self.chunk_overlap,
            }

    def index_local_file(self, file_path: str) -> Dict[str, Any]:
        """
        Index a local file from the filesystem into RAG.

        Args:
            file_path: Path to the file to index

        Returns:
            Dict with status, chunk_count, and any errors
        """
        from pathlib import Path
        import mimetypes

        file_path = Path(file_path)

        if not file_path.exists():
            return {"status": "error", "error": f"File not found: {file_path}"}

        if not file_path.is_file():
            return {"status": "error", "error": f"Not a file: {file_path}"}

        # Determine file type
        mime_type, _ = mimetypes.guess_type(str(file_path))

        # Read file content based on type
        try:
            if file_path.suffix.lower() in [".txt", ".md", ".markdown"]:
                # Text files
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()
            elif file_path.suffix.lower() in [".html", ".htm"]:
                # HTML files - strip tags
                from bs4 import BeautifulSoup

                with open(file_path, "r", encoding="utf-8") as f:
                    soup = BeautifulSoup(f.read(), "html.parser")
                    content = soup.get_text()
            elif file_path.suffix.lower() == ".pdf":
                # PDF files - extract text
                from pypdf import PdfReader

                content = ""
                with open(file_path, "rb") as f:
                    pdf_reader = PdfReader(f)
                    for page in pdf_reader.pages:
                        content += page.extract_text()
            else:
                return {
                    "status": "skipped",
                    "error": f"Unsupported file type: {file_path.suffix}",
                }

            if not content or len(content.strip()) < 10:
                return {
                    "status": "error",
                    "error": "File has no extractable text content",
                }

            # Create LangChain Document from text
            doc = LangchainDocument(
                page_content=content,
                metadata={
                    "source": str(file_path),
                    "source_id": f"local_{file_path.stem}_{hash(str(file_path))}",
                    "title": file_path.stem,
                    "document_title": file_path.stem,
                    "file_type": file_path.suffix.lower(),
                    "file_size": file_path.stat().st_size,
                    "source_type": "local_file",
                    "collection": "local_library",
                },
            )

            # Split into chunks
            chunks = self.text_splitter.split_documents([doc])
            logger.info(
                f"Split local file {file_path} into {len(chunks)} chunks"
            )

            # Store chunks in database (returns UUID-based IDs)
            embedding_ids = self.embedding_manager._store_chunks_to_db(
                chunks=chunks,
                collection_name="local_library",
                source_type="local_file",
                source_id=str(file_path),
            )

            # Load or create FAISS index using default library collection
            if self.faiss_index is None:
                from ...database.library_init import get_default_library_id

                default_collection_id = get_default_library_id(
                    self.username, self.db_password
                )
                self.faiss_index = self.load_or_create_faiss_index(
                    default_collection_id
                )

            # Filter out chunks that already exist in FAISS and deduplicate
            if self.faiss_index is not None:
                existing_ids = (
                    set(self.faiss_index.docstore._dict.keys())
                    if hasattr(self.faiss_index, "docstore")
                    else set()
                )
            else:
                existing_ids = None
            new_chunks, new_ids = self._deduplicate_chunks(
                chunks, embedding_ids, existing_ids
            )

            # Add embeddings to FAISS index
            if new_chunks:
                self.faiss_index.add_documents(new_chunks, ids=new_ids)

            # Save FAISS index
            index_path = (
                Path(self.rag_index_record.index_path)
                if self.rag_index_record
                else None
            )
            if index_path:
                self.faiss_index.save_local(
                    str(index_path.parent), index_name=index_path.stem
                )
                # Record file integrity
                self.integrity_manager.record_file(
                    index_path,
                    related_entity_type="rag_index",
                    related_entity_id=self.rag_index_record.id,
                )
                logger.info(
                    f"Saved FAISS index to {index_path} with integrity tracking"
                )

            logger.info(
                f"Successfully indexed local file {file_path} with {len(new_chunks)} new chunks "
                f"({len(chunks) - len(new_chunks)} skipped)"
            )

            return {
                "status": "success",
                "chunk_count": len(new_chunks),
                "embedding_ids": new_ids,
            }

        except Exception as e:
            logger.exception(f"Error indexing local file {file_path}")
            return {
                "status": "error",
                "error": f"Operation failed: {type(e).__name__}",
            }

    def index_user_document(
        self, user_doc, collection_name: str, force_reindex: bool = False
    ) -> Dict[str, Any]:
        """
        Index a user-uploaded document into a specific collection.

        Args:
            user_doc: UserDocument object
            collection_name: Name of the collection (e.g., "collection_123")
            force_reindex: Whether to force reindexing

        Returns:
            Dict with status, chunk_count, and any errors
        """

        try:
            # Use the pre-extracted text content
            content = user_doc.text_content

            if not content or len(content.strip()) < 10:
                return {
                    "status": "error",
                    "error": "Document has no extractable text content",
                }

            # Create LangChain Document
            doc = LangchainDocument(
                page_content=content,
                metadata={
                    "source": f"user_upload_{user_doc.id}",
                    "source_id": user_doc.id,
                    "title": user_doc.filename,
                    "document_title": user_doc.filename,
                    "file_type": user_doc.file_type,
                    "file_size": user_doc.file_size,
                    "collection": collection_name,
                },
            )

            # Split into chunks
            chunks = self.text_splitter.split_documents([doc])
            logger.info(
                f"Split user document {user_doc.filename} into {len(chunks)} chunks"
            )

            # Store chunks in database
            embedding_ids = self.embedding_manager._store_chunks_to_db(
                chunks=chunks,
                collection_name=collection_name,
                source_type="user_document",
                source_id=user_doc.id,
            )

            # Load or create FAISS index for this collection
            if self.faiss_index is None:
                # Extract collection_id from collection_name (format: "collection_<uuid>")
                collection_id = collection_name.removeprefix("collection_")
                self.faiss_index = self.load_or_create_faiss_index(
                    collection_id
                )

            # If force_reindex, remove old chunks from FAISS before adding new ones
            if force_reindex:
                existing_ids = (
                    set(self.faiss_index.docstore._dict.keys())
                    if hasattr(self.faiss_index, "docstore")
                    else set()
                )
                old_chunk_ids = list(
                    {eid for eid in embedding_ids if eid in existing_ids}
                )
                if old_chunk_ids:
                    logger.info(
                        f"Force re-index: removing {len(old_chunk_ids)} existing chunks from FAISS"
                    )
                    self.faiss_index.delete(old_chunk_ids)

            # Filter out chunks that already exist in FAISS (unless force_reindex)
            if not force_reindex:
                existing_ids = (
                    set(self.faiss_index.docstore._dict.keys())
                    if hasattr(self.faiss_index, "docstore")
                    else set()
                )
            else:
                existing_ids = None

            unique_count = len(set(embedding_ids))
            batch_dups = len(chunks) - unique_count

            new_chunks, new_ids = self._deduplicate_chunks(
                chunks, embedding_ids, existing_ids
            )

            # Add embeddings to FAISS index
            if new_chunks:
                if force_reindex:
                    logger.info(
                        f"Force re-index: adding {len(new_chunks)} chunks with updated metadata to FAISS index"
                    )
                else:
                    already_exist = unique_count - len(new_chunks)
                    logger.info(
                        f"Adding {len(new_chunks)} new chunks to FAISS index "
                        f"({already_exist} already exist, {batch_dups} batch duplicates removed)"
                    )
                self.faiss_index.add_documents(new_chunks, ids=new_ids)
            else:
                logger.info(
                    f"All {len(chunks)} chunks already exist in FAISS index, skipping"
                )

            # Save FAISS index
            index_path = (
                Path(self.rag_index_record.index_path)
                if self.rag_index_record
                else None
            )
            if index_path:
                self.faiss_index.save_local(
                    str(index_path.parent), index_name=index_path.stem
                )
                # Record file integrity
                self.integrity_manager.record_file(
                    index_path,
                    related_entity_type="rag_index",
                    related_entity_id=self.rag_index_record.id,
                )

            logger.info(
                f"Successfully indexed user document {user_doc.filename} with {len(chunks)} chunks"
            )

            return {
                "status": "success",
                "chunk_count": len(chunks),
                "embedding_ids": embedding_ids,
            }

        except Exception as e:
            logger.exception(
                f"Error indexing user document {user_doc.filename}"
            )
            return {
                "status": "error",
                "error": f"Operation failed: {type(e).__name__}",
            }

    def remove_collection_from_index(
        self, collection_name: str
    ) -> Dict[str, Any]:
        """
        Remove all documents from a collection from the FAISS index.

        Args:
            collection_name: Name of the collection (e.g., "collection_123")

        Returns:
            Dict with status and count of removed chunks
        """
        from ...database.models import DocumentChunk
        from ...database.session_context import get_user_db_session

        try:
            with get_user_db_session(
                self.username, self.db_password
            ) as session:
                # Get all chunk IDs for this collection
                chunks = (
                    session.query(DocumentChunk)
                    .filter_by(collection_name=collection_name)
                    .all()
                )

                if not chunks:
                    return {"status": "success", "deleted_count": 0}

                chunk_ids = [
                    f"{collection_name}_{chunk.id}" for chunk in chunks
                ]

                # Load FAISS index if not already loaded
                if self.faiss_index is None:
                    # Extract collection_id from collection_name (format: "collection_<uuid>")
                    collection_id = collection_name.removeprefix("collection_")
                    self.faiss_index = self.load_or_create_faiss_index(
                        collection_id
                    )

                # Remove from FAISS index
                if hasattr(self.faiss_index, "delete"):
                    try:
                        self.faiss_index.delete(chunk_ids)

                        # Save updated index
                        index_path = (
                            Path(self.rag_index_record.index_path)
                            if self.rag_index_record
                            else None
                        )
                        if index_path:
                            self.faiss_index.save_local(
                                str(index_path.parent),
                                index_name=index_path.stem,
                            )
                            # Record file integrity
                            self.integrity_manager.record_file(
                                index_path,
                                related_entity_type="rag_index",
                                related_entity_id=self.rag_index_record.id,
                            )
                    except Exception:
                        logger.warning("Could not delete chunks from FAISS")

                logger.info(
                    f"Removed {len(chunk_ids)} chunks from collection {collection_name}"
                )

                return {"status": "success", "deleted_count": len(chunk_ids)}

        except Exception as e:
            logger.exception(
                f"Error removing collection {collection_name} from index"
            )
            return {
                "status": "error",
                "error": f"Operation failed: {type(e).__name__}",
            }
