import hashlib
import json
import os
import time
import uuid
from concurrent.futures import ProcessPoolExecutor
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional

import numpy as np
from faiss import IndexFlatL2
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_community.document_loaders import (
    CSVLoader,
    PyPDFLoader,
    TextLoader,
    UnstructuredExcelLoader,
    UnstructuredMarkdownLoader,
    UnstructuredWordDocumentLoader,
)
from langchain_community.document_loaders.base import BaseLoader
from langchain_community.embeddings import (
    HuggingFaceEmbeddings,
)
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.language_models import BaseLLM
from langchain_text_splitters import RecursiveCharacterTextSplitter
from loguru import logger

from ...config import search_config
from ...database.models.library import DocumentChunk
from ...database.session_context import get_user_db_session
from ...utilities.url_utils import normalize_url
from ..search_engine_base import BaseSearchEngine


def _get_file_loader(file_path: str) -> Optional[BaseLoader]:
    """Get an appropriate document loader for a file based on its extension"""
    file_path = Path(file_path)
    extension = file_path.suffix.lower()

    try:
        if extension == ".pdf":
            return PyPDFLoader(str(file_path))
        elif extension == ".txt":
            return TextLoader(str(file_path))
        elif extension in [".md", ".markdown"]:
            return UnstructuredMarkdownLoader(str(file_path))
        elif extension in [".doc", ".docx"]:
            return UnstructuredWordDocumentLoader(str(file_path))
        elif extension == ".csv":
            return CSVLoader(str(file_path))
        elif extension in [".xls", ".xlsx"]:
            return UnstructuredExcelLoader(str(file_path))
        else:
            # Try the text loader as a fallback for unknown extensions
            logger.warning(
                f"Unknown file extension for {file_path}, trying TextLoader"
            )
            return TextLoader(str(file_path), encoding="utf-8")
    except Exception:
        logger.exception(f"Error creating loader for {file_path}")
        return None


def _load_document(file_path: Path) -> List[Document]:
    """
    Loads documents from a file.

    Args:
        file_path: The path to the document to load.

    Returns:
        The loaded documents, or an empty list if it failed to load.

    """
    # Get a loader for this file
    loader = _get_file_loader(str(file_path))

    if loader is None:
        # No loader for this filetype.
        return []

    try:
        # Load the document
        docs = loader.load()

        # Add source path metadata and ID.
        for doc in docs:
            doc.metadata["source"] = str(file_path)
            doc.metadata["filename"] = file_path.name

    except Exception:
        logger.exception(f"Error loading {file_path}")
        return []

    return docs


class LocalEmbeddingManager:
    """Handles embedding generation and storage for local document search"""

    def __init__(
        self,
        embedding_model: str = "all-MiniLM-L6-v2",
        embedding_device: str = "cpu",
        embedding_model_type: str = "sentence_transformers",  # or 'ollama'
        ollama_base_url: Optional[str] = None,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        cache_dir: str = ".cache/local_search",
        settings_snapshot: Optional[Dict[str, Any]] = None,
    ):
        """
        Initialize the embedding manager for local document search.

        Args:
            embedding_model: Name of the embedding model to use
            embedding_device: Device to run embeddings on ('cpu' or 'cuda')
            embedding_model_type: Type of embedding model ('sentence_transformers' or 'ollama')
            ollama_base_url: Base URL for Ollama API if using ollama embeddings
            chunk_size: Size of text chunks for splitting documents
            chunk_overlap: Overlap between chunks
            cache_dir: Directory to store embedding cache and index
            settings_snapshot: Optional settings snapshot for background threads
        """

        self.embedding_model = embedding_model
        self.embedding_device = embedding_device
        self.embedding_model_type = embedding_model_type
        self.ollama_base_url = ollama_base_url
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.cache_dir = Path(cache_dir)
        self.settings_snapshot = settings_snapshot or {}

        # Username for database access (extracted from settings if available)
        self.username = (
            settings_snapshot.get("_username") if settings_snapshot else None
        )
        # Password for encrypted database access (can be set later)
        self.db_password = None

        # Create cache directory if it doesn't exist
        self.cache_dir.mkdir(parents=True, exist_ok=True)

        # Initialize the embedding model
        self._embeddings = None

        # Initialize the text splitter
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size, chunk_overlap=chunk_overlap
        )

        # Track indexed folders and their metadata
        self.indexed_folders = self._load_indexed_folders()

        # Vector store cache
        self.vector_stores = {}

    @property
    def embeddings(self):
        """
        Lazily initialize embeddings when first accessed.
        This allows the LocalEmbeddingManager to be created without
        immediately loading models, which is helpful when no local search is performed.
        """
        if self._embeddings is None:
            logger.info("Initializing embeddings on first use")
            self._embeddings = self._initialize_embeddings()
        return self._embeddings

    def _initialize_embeddings(self):
        """Initialize the embedding model based on configuration"""
        try:
            # Use the new unified embedding system
            from ...embeddings import get_embeddings

            # Prepare kwargs for provider-specific parameters
            kwargs = {}

            # Add device for sentence transformers
            if self.embedding_model_type == "sentence_transformers":
                kwargs["device"] = self.embedding_device

            # Add base_url for ollama if specified
            if self.embedding_model_type == "ollama" and self.ollama_base_url:
                kwargs["base_url"] = normalize_url(self.ollama_base_url)

            logger.info(
                f"Initializing embeddings with provider={self.embedding_model_type}, model={self.embedding_model}"
            )

            return get_embeddings(
                provider=self.embedding_model_type,
                model=self.embedding_model,
                settings_snapshot=self.settings_snapshot,
                **kwargs,
            )
        except Exception:
            logger.exception("Error initializing embeddings")
            logger.warning(
                "Falling back to HuggingFaceEmbeddings with all-MiniLM-L6-v2"
            )
            return HuggingFaceEmbeddings(
                model_name="sentence-transformers/all-MiniLM-L6-v2"
            )

    def _store_chunks_to_db(
        self,
        chunks: List[Document],
        collection_name: str,
        source_path: Optional[str] = None,
        source_id: Optional[int] = None,
        source_type: str = "local_file",
    ) -> List[str]:
        """
        Store document chunks in the database.

        Args:
            chunks: List of LangChain Document chunks
            collection_name: Name of the collection (e.g., 'personal_notes', 'library')
            source_path: Path to source file (for local files)
            source_id: ID of source document (for library documents)
            source_type: Type of source ('local_file' or 'library')

        Returns:
            List of chunk embedding IDs (UUIDs) for FAISS mapping
        """
        if not self.username:
            logger.warning(
                "No username available, cannot store chunks in database"
            )
            return []

        chunk_ids = []

        try:
            with get_user_db_session(
                self.username, self.db_password
            ) as session:
                for idx, chunk in enumerate(chunks):
                    # Generate unique hash for chunk
                    chunk_text = chunk.page_content
                    chunk_hash = hashlib.sha256(chunk_text.encode()).hexdigest()

                    # Generate unique embedding ID
                    embedding_id = uuid.uuid4().hex

                    # Extract metadata
                    metadata = chunk.metadata or {}
                    document_title = metadata.get(
                        "filename", metadata.get("title", "Unknown")
                    )

                    # Calculate word count
                    word_count = len(chunk_text.split())

                    # Get character positions from metadata if available
                    start_char = metadata.get("start_char", 0)
                    end_char = metadata.get("end_char", len(chunk_text))

                    # Check if chunk already exists
                    existing_chunk = (
                        session.query(DocumentChunk)
                        .filter_by(chunk_hash=chunk_hash)
                        .first()
                    )

                    if existing_chunk:
                        # Update existing chunk
                        existing_chunk.last_accessed = datetime.now(UTC)
                        chunk_ids.append(existing_chunk.embedding_id)
                        logger.debug(
                            f"Chunk already exists, reusing: {existing_chunk.embedding_id}"
                        )
                    else:
                        # Create new chunk
                        db_chunk = DocumentChunk(
                            chunk_hash=chunk_hash,
                            source_type=source_type,
                            source_id=source_id,
                            source_path=str(source_path)
                            if source_path
                            else None,
                            collection_name=collection_name,
                            chunk_text=chunk_text,
                            chunk_index=idx,
                            start_char=start_char,
                            end_char=end_char,
                            word_count=word_count,
                            embedding_id=embedding_id,
                            embedding_model=self.embedding_model,
                            embedding_model_type=self.embedding_model_type,
                            document_title=document_title,
                            document_metadata=metadata,
                        )
                        session.add(db_chunk)
                        chunk_ids.append(embedding_id)

                session.commit()
                logger.info(
                    f"Stored {len(chunk_ids)} chunks to database for collection '{collection_name}'"
                )

        except Exception:
            logger.exception(
                f"Error storing chunks to database for collection '{collection_name}'"
            )
            return []

        return chunk_ids

    def _load_chunks_from_db(
        self, chunk_ids: List[str], username: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Load chunks from database by their embedding IDs.

        Args:
            chunk_ids: List of embedding IDs to load
            username: Username for database access (uses self.username if not provided)

        Returns:
            List of chunk dictionaries with content and metadata
        """
        username = username or self.username
        if not username:
            logger.warning(
                "No username available, cannot load chunks from database"
            )
            return []

        chunks = []

        try:
            with get_user_db_session(username) as session:
                db_chunks = (
                    session.query(DocumentChunk)
                    .filter(DocumentChunk.embedding_id.in_(chunk_ids))
                    .all()
                )

                for db_chunk in db_chunks:
                    # Update last accessed time
                    db_chunk.last_accessed = datetime.now(UTC)

                    chunks.append(
                        {
                            "id": db_chunk.embedding_id,
                            "content": db_chunk.chunk_text,
                            "metadata": {
                                "source_type": db_chunk.source_type,
                                "source_path": db_chunk.source_path,
                                "source_id": db_chunk.source_id,
                                "collection": db_chunk.collection_name,
                                "chunk_index": db_chunk.chunk_index,
                                "word_count": db_chunk.word_count,
                                "title": db_chunk.document_title,
                                **db_chunk.document_metadata,
                            },
                        }
                    )

                session.commit()  # Commit the last_accessed updates

        except Exception:
            logger.exception("Error loading chunks from database")
            return []

        return chunks

    def _delete_chunks_from_db(
        self,
        collection_name: str,
        source_path: Optional[str] = None,
        source_id: Optional[int] = None,
    ) -> int:
        """
        Delete chunks from database.

        Args:
            collection_name: Name of the collection
            source_path: Path to source file (for local files)
            source_id: ID of source document (for library documents)

        Returns:
            Number of chunks deleted
        """
        if not self.username:
            logger.warning(
                "No username available, cannot delete chunks from database"
            )
            return 0

        try:
            with get_user_db_session(
                self.username, self.db_password
            ) as session:
                query = session.query(DocumentChunk).filter_by(
                    collection_name=collection_name
                )

                if source_path:
                    query = query.filter_by(source_path=str(source_path))
                if source_id:
                    query = query.filter_by(source_id=source_id)

                count = query.delete()
                session.commit()

                logger.info(
                    f"Deleted {count} chunks from database for collection '{collection_name}'"
                )
                return count

        except Exception:
            logger.exception(
                f"Error deleting chunks from database for collection '{collection_name}'"
            )
            return 0

    def _load_or_create_vector_store(self):
        """Load the vector store from disk or create it if needed"""
        vector_store_path = self._get_vector_store_path()

        # Check if vector store exists and is up to date
        if vector_store_path.exists() and not self._check_folders_modified():
            logger.info(
                f"Loading existing vector store from {vector_store_path}"
            )
            try:
                vector_store = FAISS.load_local(
                    str(vector_store_path),
                    self.embeddings,
                    allow_dangerous_deserialization=True,
                )

                # Add this code to show document count
                doc_count = len(vector_store.index_to_docstore_id)
                logger.info(f"Loaded index with {doc_count} document chunks")

                return vector_store
            except Exception:
                logger.exception("Error loading vector store")
                logger.info("Will create a new vector store")

        # Create a new vector store
        return self._create_vector_store()

    def _load_indexed_folders(self) -> Dict[str, Dict[str, Any]]:
        """Load metadata about indexed folders from disk"""
        index_metadata_path = self.cache_dir / "index_metadata.json"

        if index_metadata_path.exists():
            try:
                with open(index_metadata_path, "r") as f:
                    return json.load(f)
            except Exception:
                logger.exception("Error loading index metadata")

        return {}

    def _save_indexed_folders(self):
        """Save metadata about indexed folders to disk"""
        index_metadata_path = self.cache_dir / "index_metadata.json"

        try:
            with open(index_metadata_path, "w") as f:
                json.dump(self.indexed_folders, f, indent=2)
        except Exception:
            logger.exception("Error saving index metadata")

    @staticmethod
    def get_folder_hash(folder_path: Path) -> str:
        """Generate a hash for a folder based on its path"""
        # Canonicalize the path so we don't have weird Windows vs. Linux
        # problems or issues with trailing slashes.
        canonical_folder_path = "/".join(folder_path.parts)
        return hashlib.md5(  # DevSkim: ignore DS126858
            canonical_folder_path.encode(), usedforsecurity=False
        ).hexdigest()

    def _get_index_path(self, folder_path: Path) -> Path:
        """Get the path where the index for a specific folder should be stored"""
        folder_hash = self.get_folder_hash(folder_path)
        return self.cache_dir / f"index_{folder_hash}"

    def _check_folder_modified(self, folder_path: Path) -> bool:
        """Check if a folder has been modified since it was last indexed"""

    @staticmethod
    def _get_all_files(folder_path: Path) -> Iterable[Path]:
        """
        Gets all the files, recursively, in a folder.

        Args:
            folder_path: The path to the folder.

        Yields:
            Each of the files in the folder.

        """
        for root, _, files in os.walk(folder_path):
            for file in files:
                yield Path(root) / file

    def _get_modified_files(self, folder_path: Path) -> List[Path]:
        """
        Gets the files in a folder that have been modified since it was last
        indexed.

        Args:
            folder_path: The path to the folder to check.

        Returns:
            A list of the files that were modified.

        """
        if not folder_path.exists() or not folder_path.is_dir():
            return []

        folder_hash = self.get_folder_hash(folder_path)

        if folder_hash not in self.indexed_folders:
            # If folder has never been indexed, everything has been modified.
            last_indexed = 0
            indexed_files = set()
        else:
            last_indexed = self.indexed_folders[folder_hash].get(
                "last_indexed", 0
            )
            indexed_files = (
                self.indexed_folders[folder_hash]
                .get("indexed_files", {})
                .keys()
            )

        # Check if any file in the folder has been modified since last indexing
        modified_files = []
        for file_path in self._get_all_files(folder_path):
            file_stats = file_path.stat()
            if file_stats.st_mtime > last_indexed:
                modified_files.append(file_path)
            elif str(file_path.relative_to(folder_path)) not in indexed_files:
                # This file somehow never got indexed.
                modified_files.append(file_path)

        return modified_files

    def _check_config_changed(self, folder_path: Path) -> bool:
        """
        Checks if the embedding configuration for a folder has been changed
        since it was last indexed.
        """
        folder_hash = self.get_folder_hash(folder_path)

        if folder_hash not in self.indexed_folders:
            # It hasn't been indexed at all. That's a new configuration,
            # technically.
            return True

        embedding_config = self.indexed_folders[folder_hash]
        chunk_size = int(embedding_config.get("chunk_size", 0))
        chunk_overlap = int(embedding_config.get("chunk_overlap", 0))
        embedding_model = embedding_config.get("embedding_model", "")

        if (chunk_size, chunk_overlap, embedding_model) != (
            self.chunk_size,
            self.chunk_overlap,
            self.embedding_model,
        ):
            logger.info(
                "Embedding configuration has changed, re-indexing folder."
            )
            return True
        return False

    def index_folder(
        self, folder_path: str, force_reindex: bool = False
    ) -> bool:
        """
        Index all documents in a folder for vector search.

        Args:
            folder_path: Path to the folder to index
            force_reindex: Whether to force reindexing even if unchanged

        Returns:
            bool: True if indexing was successful, False otherwise
        """
        folder_path = Path(folder_path)

        # Validate folder
        if not folder_path.exists():
            logger.error(f"Folder not found: {folder_path}")
            return False

        if not folder_path.is_dir():
            logger.error(f"Path is not a directory: {folder_path}")
            return False

        folder_str = str(folder_path)
        folder_hash = self.get_folder_hash(folder_path)
        index_path = self._get_index_path(folder_path)

        if force_reindex or self._check_config_changed(folder_path):
            logger.info(f"Re-indexing entire folder: {folder_path}")
            modified_files = list(self._get_all_files(folder_path))
        else:
            # Just re-index the modified files if we can get away with it.
            modified_files = self._get_modified_files(folder_path)
            logger.info(f"Re-indexing {len(modified_files)} modified files...")

        # Load the vector store from disk if not already loaded
        if folder_hash not in self.vector_stores and index_path.exists():
            try:
                self.vector_stores[folder_hash] = FAISS.load_local(
                    str(index_path),
                    self.embeddings,
                    allow_dangerous_deserialization=True,
                )
                logger.info(f"Loaded index for {folder_path} from disk")
            except Exception:
                logger.exception(f"Error loading index for {folder_path}")
                # If loading fails, force reindexing
                force_reindex = True

        logger.info(f"Indexing folder: {folder_path}")
        start_time = time.time()

        # Find documents to index
        all_docs = []

        # Remove hidden files and directories.
        modified_files = [
            p
            for p in modified_files
            if not p.name.startswith(".")
            and not any(part.startswith(".") for part in p.parts)
        ]
        # Index them.
        with ProcessPoolExecutor() as executor:
            all_docs_nested = executor.map(_load_document, modified_files)
        # Flatten the result.
        for docs in all_docs_nested:
            all_docs.extend(docs)

        if force_reindex or folder_hash not in self.vector_stores:
            logger.info(f"Creating new index for {folder_path}")
            # Embed a test query to figure out embedding length.
            test_embedding = self.embeddings.embed_query("hello world")
            index = IndexFlatL2(len(test_embedding))
            # Use minimal docstore - chunks are stored in database
            self.vector_stores[folder_hash] = FAISS(
                self.embeddings,
                index=index,
                docstore=InMemoryDocstore(),  # Minimal - just for FAISS compatibility
                index_to_docstore_id={},
                normalize_L2=True,
            )

        # Split documents into chunks
        logger.info(f"Splitting {len(all_docs)} documents into chunks")
        splits = self.text_splitter.split_documents(all_docs)
        logger.info(
            f"Created {len(splits)} chunks from {len(modified_files)} files"
        )

        # Store chunks in database and get embedding IDs
        embedding_ids = []
        if splits:
            logger.info(f"Storing {len(splits)} chunks in database")
            # Get collection name from folder path (last folder name)
            collection_name = folder_path.name

            # Store chunks to database
            embedding_ids = self._store_chunks_to_db(
                chunks=splits,
                collection_name=collection_name,
                source_type="local_file",
            )

            logger.info(f"Adding {len(splits)} chunks to FAISS index")
            # Add embeddings to FAISS using the database-generated IDs
            self.vector_stores[folder_hash].add_documents(
                splits, ids=embedding_ids
            )

        # Update indexing time for individual files.
        index_time = time.time()
        indexed_files = {}
        if folder_hash in self.indexed_folders:
            indexed_files = (
                self.indexed_folders[folder_hash]
                .get("indexed_files", {})
                .copy()
            )
        for embedding_id, split in zip(embedding_ids, splits, strict=False):
            split_source = str(
                Path(split.metadata["source"]).relative_to(folder_path)
            )
            id_list = indexed_files.setdefault(split_source, [])
            id_list.append(embedding_id)

        # Check for any files that were removed and remove them from the
        # vector store and database.
        delete_ids = []
        delete_paths = []
        for relative_path, chunk_ids in indexed_files.items():
            if not (folder_path / Path(relative_path)).exists():
                delete_ids.extend(chunk_ids)
                delete_paths.append(relative_path)
        if delete_ids:
            logger.info(
                f"Deleting {len(delete_paths)} non-existent files from the "
                f"index and database."
            )
            # Delete from FAISS index
            self.vector_stores[folder_hash].delete(delete_ids)

            # Delete from database
            collection_name = folder_path.name
            for delete_path in delete_paths:
                full_path = str(folder_path / delete_path)
                deleted_count = self._delete_chunks_from_db(
                    collection_name=collection_name,
                    source_path=full_path,
                )
                logger.debug(
                    f"Deleted {deleted_count} chunks for {delete_path} from database"
                )
        for path in delete_paths:
            del indexed_files[path]

        # Save the vector store to disk
        logger.info(f"Saving index to {index_path}")
        self.vector_stores[folder_hash].save_local(str(index_path))

        # Update metadata
        self.indexed_folders[folder_hash] = {
            "path": folder_str,
            "last_indexed": index_time,
            "file_count": len(modified_files),
            "chunk_count": len(splits),
            "embedding_model": self.embedding_model,
            "chunk_size": self.chunk_size,
            "chunk_overlap": self.chunk_overlap,
            "indexed_files": indexed_files,
        }

        # Save updated metadata
        self._save_indexed_folders()

        elapsed_time = time.time() - start_time
        logger.info(
            f"Indexed {len(modified_files)} files in {elapsed_time:.2f} seconds"
        )

        return True

    def search(
        self,
        query: str,
        folder_paths: List[str],
        limit: int = 10,
        score_threshold: float = 0.0,
    ) -> List[Dict[str, Any]]:
        """
        Search for documents relevant to a query across specified folders.

        Args:
            query: The search query
            folder_paths: List of folder paths to search in
            limit: Maximum number of results to return
            score_threshold: Minimum similarity score threshold

        Returns:
            List of results with document content and metadata
        """
        folder_paths = [Path(p) for p in folder_paths]

        # Add detailed debugging for each folder
        for folder_path in folder_paths:
            folder_hash = self.get_folder_hash(folder_path)
            index_path = self._get_index_path(folder_path)

            logger.info(f"Diagnostic for {folder_path}:")
            logger.info(f"  - Folder hash: {folder_hash}")
            logger.info(f"  - Index path: {index_path}")
            logger.info(f"  - Index exists on disk: {index_path.exists()}")
            logger.info(
                f"  - Is in indexed_folders: {folder_hash in self.indexed_folders}"
            )

            if folder_hash in self.indexed_folders:
                meta = self.indexed_folders[folder_hash]
                logger.info(
                    f"  - Metadata: file_count={meta.get('file_count', 0)}, chunk_count={meta.get('chunk_count', 0)}"
                )

        # Validate folders exist
        valid_folder_paths = []
        for path in folder_paths:
            if path.exists() and path.is_dir():
                valid_folder_paths.append(path)
            else:
                logger.warning(
                    f"Skipping non-existent folder in search: {path}"
                )

        # If no valid folders, return empty results
        if not valid_folder_paths:
            logger.warning(f"No valid folders to search among: {folder_paths}")
            return []

        all_results = []

        for folder_path in valid_folder_paths:
            folder_hash = self.get_folder_hash(folder_path)

            # Skip folders that haven't been indexed
            if folder_hash not in self.indexed_folders:
                logger.warning(f"Folder {folder_path} has not been indexed")
                continue

            # Make sure the vector store is loaded
            if folder_hash not in self.vector_stores:
                index_path = self._get_index_path(folder_path)
                try:
                    self.vector_stores[folder_hash] = FAISS.load_local(
                        str(index_path),
                        self.embeddings,
                        allow_dangerous_deserialization=True,
                    )
                except Exception:
                    logger.exception(f"Error loading index for {folder_path}")
                    continue

            # Search in this folder
            vector_store = self.vector_stores[folder_hash]

            try:
                # Get query embedding
                query_vector = self.embeddings.embed_query(query)

                # Search FAISS index for similar vectors
                # Returns: (distances, indices) where indices are FAISS internal indices
                distances, indices = vector_store.index.search(
                    np.array([query_vector], dtype=np.float32), limit
                )

                # Convert distances to similarity scores (L2 distance -> similarity)
                # For L2: smaller distance = more similar
                # Convert to similarity: 1 / (1 + distance)
                similarities = 1 / (1 + distances[0])

                # Get embedding IDs from FAISS mapping
                embedding_ids = []
                valid_indices = []
                for idx, faiss_idx in enumerate(indices[0]):
                    if faiss_idx == -1:  # FAISS returns -1 for empty results
                        continue
                    if faiss_idx in vector_store.index_to_docstore_id:
                        embedding_id = vector_store.index_to_docstore_id[
                            faiss_idx
                        ]
                        embedding_ids.append(embedding_id)
                        valid_indices.append(idx)

                # Load chunks from database
                if embedding_ids:
                    db_chunks = self._load_chunks_from_db(
                        embedding_ids, self.username
                    )

                    # Create results from database chunks
                    for idx, chunk in zip(valid_indices, db_chunks):
                        similarity = float(similarities[idx])

                        # Skip results below the threshold
                        if similarity < score_threshold:
                            continue

                        # Extract metadata from chunk
                        metadata = chunk.get("document_metadata", {})
                        if "source" not in metadata and chunk.get(
                            "source_path"
                        ):
                            metadata["source"] = chunk["source_path"]

                        result = {
                            "content": chunk["chunk_text"],
                            "metadata": metadata,
                            "similarity": similarity,
                            "folder": folder_path,
                        }

                        all_results.append(result)
            except Exception:
                logger.exception(f"Error searching in {folder_path}")

        # Sort by similarity (highest first)
        all_results.sort(key=lambda x: x["similarity"], reverse=True)

        # Limit to the requested number
        return all_results[:limit]

    def clear_cache(self):
        """Clear all cached vector stores from memory (not disk)"""
        self.vector_stores.clear()

    def get_indexed_folders_info(self) -> List[Dict[str, Any]]:
        """Get information about all indexed folders"""
        info = []

        for folder_hash, metadata in self.indexed_folders.items():
            folder_info = metadata.copy()

            # Add formatted last indexed time
            if "last_indexed" in folder_info:
                folder_info["last_indexed_formatted"] = datetime.fromtimestamp(
                    folder_info["last_indexed"]
                ).strftime("%Y-%m-%d %H:%M:%S")

            # Check if index file exists
            index_path = self._get_index_path(Path(folder_info["path"]))
            folder_info["index_exists"] = index_path.exists()

            info.append(folder_info)

        return info


class LocalSearchEngine(BaseSearchEngine):
    """Local document search engine with two-phase retrieval"""

    def __init__(
        self,
        paths: List[str],
        llm: Optional[BaseLLM] = None,
        max_results: int = 10,
        max_filtered_results: Optional[int] = None,
        embedding_model: str = "all-MiniLM-L6-v2",
        embedding_device: str = "cpu",
        embedding_model_type: str = "sentence_transformers",
        ollama_base_url: Optional[str] = None,
        force_reindex: bool = False,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        cache_dir: str = ".cache/local_search",
        collections: Optional[Dict[str, Dict[str, Any]]] = None,
        name: str = "",
        description: str = "",
    ):
        """
        Initialize the local search engine.

        Args:
            paths: List of folder paths to search in
            llm: Language model for relevance filtering
            max_results: Maximum number of results to return
            max_filtered_results: Maximum results after filtering
            embedding_model: Name of the embedding model to use
            embedding_device: Device to run embeddings on ('cpu' or 'cuda')
            embedding_model_type: Type of embedding model
            ollama_base_url: Base URL for Ollama API
            force_reindex: Whether to force reindexing
            chunk_size: Size of text chunks for splitting documents
            chunk_overlap: Overlap between chunks
            cache_dir: Directory to store embedding cache and index
            collections: Dictionary of named collections with paths and descriptions
            name: Human-readable name of the collection we are searching.
            description: Human-readable description of the collection we are
                searching.
        """
        # Initialize the base search engine
        super().__init__(llm=llm, max_filtered_results=max_filtered_results)

        self.name = name
        self.description = description

        # Validate folder paths
        self.folder_paths = paths
        self.valid_folder_paths = []
        for path_str in paths:
            path = Path(path_str)
            if path.exists() and path.is_dir():
                self.valid_folder_paths.append(path_str)
            else:
                logger.warning(
                    f"Folder not found or is not a directory: {path_str}"
                )

        # If no valid folders, log a clear message
        if not self.valid_folder_paths and paths:
            logger.warning(f"No valid folders found among: {paths}")
            logger.warning(
                "This search engine will return no results until valid folders are configured"
            )

        self.max_results = max_results
        self.collections = collections or {
            "default": {"paths": paths, "description": "Default collection"}
        }

        # Initialize the embedding manager with only valid folders
        self.embedding_manager = LocalEmbeddingManager(
            embedding_model=embedding_model,
            embedding_device=embedding_device,
            embedding_model_type=embedding_model_type,
            ollama_base_url=ollama_base_url,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            cache_dir=cache_dir,
            settings_snapshot=self.settings_snapshot,
        )

        # Index all folders
        self._index_folders(force_reindex)

    def _index_folders(self, force_reindex: bool = False):
        """Index all valid configured folders"""
        indexed = []
        failed = []
        skipped = []

        # Keep track of invalid folders
        for folder in self.folder_paths:
            if folder not in self.valid_folder_paths:
                skipped.append(folder)
                continue

            success = self.embedding_manager.index_folder(folder, force_reindex)
            if success:
                indexed.append(folder)
            else:
                failed.append(folder)

        if indexed:
            logger.info(
                f"Successfully indexed {len(indexed)} folders: {', '.join(indexed)}"
            )

        if failed:
            logger.warning(
                f"Failed to index {len(failed)} folders: {', '.join(failed)}"
            )

        if skipped:
            logger.warning(
                f"Skipped {len(skipped)} invalid folders: {', '.join(skipped)}"
            )

    def _get_previews(
        self, query: str, collection_names: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """
        Get preview information for documents matching the query.

        Args:
            query: The search query
            collection_names: Specific collections to search within (if None, search all)

        Returns:
            List of preview dictionaries
        """
        # Determine which collections to search
        if collection_names:
            # Search only in specified collections
            collections_to_search = {
                name: self.collections[name]
                for name in collection_names
                if name in self.collections
            }
            if not collections_to_search:
                logger.warning(
                    f"No valid collections found among: {collection_names}"
                )
                return []
        else:
            # Search in all collections
            collections_to_search = self.collections

        # Extract all folder paths from the collections to search
        search_paths = []
        for collection_config in collections_to_search.values():
            if "paths" in collection_config:
                search_paths.extend(collection_config["paths"])

        logger.info(
            f"Searching local documents in collections: {list(collections_to_search.keys())}"
        )

        # Filter out invalid paths
        valid_search_paths = [
            path for path in search_paths if path in self.valid_folder_paths
        ]

        if not valid_search_paths:
            logger.warning(
                f"No valid folders to search in collections: {list(collections_to_search.keys())}"
            )
            return []

        # Search across the valid selected folders
        raw_results = self.embedding_manager.search(
            query=query,
            folder_paths=valid_search_paths,
            limit=self.max_results,
            score_threshold=0.1,  # Skip very low relevance results
        )

        if not raw_results:
            logger.info(f"No local documents found for query: {query}")
            return []

        # Convert to preview format
        previews = []
        for i, result in enumerate(raw_results):
            # Create a unique ID
            result_id = f"local-{i}-{hashlib.md5(result['content'][:50].encode(), usedforsecurity=False).hexdigest()}"  # DevSkim: ignore DS126858

            # Extract filename and path
            source_path = result["metadata"].get("source", "Unknown")
            filename = result["metadata"].get(
                "filename", Path(source_path).name
            )

            # Create preview snippet (first ~200 chars of content)
            snippet = (
                result["content"][:200] + "..."
                if len(result["content"]) > 200
                else result["content"]
            )

            # Determine which collection this document belongs to
            collection_name = "Unknown"
            folder_path = result["folder"]
            for name, collection in self.collections.items():
                if any(
                    folder_path.is_relative_to(path)
                    for path in collection.get("paths", [])
                ):
                    break

            # Format the preview
            preview = {
                "id": result_id,
                "title": filename,
                "snippet": snippet,
                "link": source_path,
                "similarity": result["similarity"],
                "folder": folder_path.as_posix(),
                "collection": collection_name,
                "collection_description": self.collections.get(
                    collection_name, {}
                ).get("description", ""),
                "_full_content": result[
                    "content"
                ],  # Store full content for later
                "_metadata": result["metadata"],  # Store metadata for later
            }

            previews.append(preview)

        logger.info(f"Found {len(previews)} local document matches")
        return previews

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant documents.
        For local search, the full content is already available.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content
        """
        # Check if we should add full content
        if (
            hasattr(search_config, "SEARCH_SNIPPETS_ONLY")
            and search_config.SEARCH_SNIPPETS_ONLY
        ):
            logger.info("Snippet-only mode, skipping full content addition")
            return relevant_items

        # For local search, we already have the full content
        results = []
        for item in relevant_items:
            # Create a copy with full content
            result = item.copy()

            # Add full content if we have it
            if "_full_content" in item:
                result["content"] = item["_full_content"]
                result["full_content"] = item["_full_content"]

                # Remove temporary fields
                if "_full_content" in result:
                    del result["_full_content"]

            # Add metadata if we have it
            if "_metadata" in item:
                result["document_metadata"] = item["_metadata"]

                # Remove temporary fields
                if "_metadata" in result:
                    del result["_metadata"]

            results.append(result)

        return results

    def run(
        self,
        query: str,
        research_context: Dict[str, Any] | None = None,
        collection_names: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Execute a search using the two-phase approach.

        Args:
            query: The search query
            research_context: Context from previous research to use.
            collection_names: Specific collections to search within (if None, search all)

        Returns:
            List of search result dictionaries with full content
        """
        logger.info("---Execute a search using Local Documents---")

        # Check if we have any special collection parameters in the query
        collection_prefix = "collection:"
        remaining_query = query
        specified_collections = []

        # Parse query for collection specifications like "collection:research_papers query terms"
        query_parts = query.split()
        for part in query_parts:
            if part.lower().startswith(collection_prefix):
                collection_name = part[len(collection_prefix) :].strip()
                if collection_name in self.collections:
                    specified_collections.append(collection_name)
                    # Remove this part from the query
                    remaining_query = remaining_query.replace(
                        part, "", 1
                    ).strip()

        # If collections were specified in the query, they override the parameter
        if specified_collections:
            collection_names = specified_collections
            query = remaining_query

        # Phase 1: Get previews (with collection filtering)
        previews = self._get_previews(query, collection_names)

        if not previews:
            return []

        # Phase 2: Filter for relevance
        relevant_items = self._filter_for_relevance(previews, query)

        if not relevant_items:
            return []

        # Phase 3: Get full content for relevant items
        if (
            hasattr(search_config, "SEARCH_SNIPPETS_ONLY")
            and search_config.SEARCH_SNIPPETS_ONLY
        ):
            logger.info("Returning snippet-only results as per config")
            results = relevant_items
        else:
            results = self._get_full_content(relevant_items)

        # Clean up temporary data
        self.embedding_manager.clear_cache()

        return results

    def get_collections_info(self) -> List[Dict[str, Any]]:
        """
        Get information about all collections, including indexing status.

        Returns:
            List of collection information dictionaries
        """
        collections_info = []

        for name, collection in self.collections.items():
            paths = collection.get("paths", [])
            paths = [Path(p) for p in paths]
            description = collection.get("description", "")

            # Get indexing information for each path
            paths_info = []
            for path in paths:
                # Check if folder exists
                exists = path.exists() and path.is_dir()

                # Check if folder is indexed
                folder_hash = self.embedding_manager.get_folder_hash(path)
                indexed = folder_hash in self.embedding_manager.indexed_folders

                # Get index details if available
                index_info = {}
                if indexed:
                    index_info = self.embedding_manager.indexed_folders[
                        folder_hash
                    ].copy()

                paths_info.append(
                    {
                        "path": path,
                        "exists": exists,
                        "indexed": indexed,
                        "index_info": index_info,
                    }
                )

            collections_info.append(
                {
                    "name": name,
                    "description": description,
                    "paths": paths,
                    "paths_info": paths_info,
                    "document_count": sum(
                        info.get("index_info", {}).get("file_count", 0)
                        for info in paths_info
                    ),
                    "chunk_count": sum(
                        info.get("index_info", {}).get("chunk_count", 0)
                        for info in paths_info
                    ),
                    "all_indexed": all(
                        info["indexed"] for info in paths_info if info["exists"]
                    ),
                }
            )

        return collections_info

    def reindex_collection(self, collection_name: str) -> bool:
        """
        Reindex a specific collection.

        Args:
            collection_name: Name of the collection to reindex

        Returns:
            True if reindexing was successful, False otherwise
        """
        if collection_name not in self.collections:
            logger.error(f"Collection '{collection_name}' not found")
            return False

        paths = self.collections[collection_name].get("paths", [])
        success = True

        for path in paths:
            if not self.embedding_manager.index_folder(
                path, force_reindex=True
            ):
                success = False

        return success

    @classmethod
    def from_config(
        cls, config_dict: Dict[str, Any], llm: Optional[BaseLLM] = None
    ) -> "LocalSearchEngine":
        """
        Create a LocalSearchEngine instance from a configuration dictionary.

        Args:
            config_dict: Configuration dictionary
            llm: Language model for relevance filtering

        Returns:
            Initialized LocalSearchEngine instance
        """
        # Required parameters
        folder_paths = []
        collections = config_dict.get("collections", {})

        # Extract all folder paths from collections
        for collection_config in collections.values():
            if "paths" in collection_config:
                folder_paths.extend(collection_config["paths"])

        # Fall back to folder_paths if no collections defined
        if not folder_paths:
            folder_paths = config_dict.get("folder_paths", [])
            # Create a default collection if using folder_paths
            if folder_paths:
                collections = {
                    "default": {
                        "paths": folder_paths,
                        "description": "Default collection",
                    }
                }

        # Optional parameters with defaults
        max_results = config_dict.get("max_results", 10)
        max_filtered_results = config_dict.get("max_filtered_results")
        embedding_model = config_dict.get("embedding_model", "all-MiniLM-L6-v2")
        embedding_device = config_dict.get("embedding_device", "cpu")
        embedding_model_type = config_dict.get(
            "embedding_model_type", "sentence_transformers"
        )
        ollama_base_url = config_dict.get("ollama_base_url")
        force_reindex = config_dict.get("force_reindex", False)
        chunk_size = config_dict.get("chunk_size", 1000)
        chunk_overlap = config_dict.get("chunk_overlap", 200)
        cache_dir = config_dict.get("cache_dir", ".cache/local_search")

        return cls(
            paths=folder_paths,
            collections=collections,
            llm=llm,
            max_results=max_results,
            max_filtered_results=max_filtered_results,
            embedding_model=embedding_model,
            embedding_device=embedding_device,
            embedding_model_type=embedding_model_type,
            ollama_base_url=ollama_base_url,
            force_reindex=force_reindex,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            cache_dir=cache_dir,
        )
