"""
FAISSStore: In-memory Vector Store implementation using FAISS.

Provides high-performance vector similarity search with:
- In-memory vector storage with FAISS indexes
- Multiple distance metrics (Cosine, L2, Inner Product)
- CPU-only execution (GPU support removed)
- MMR (Maximal Marginal Relevance) search
- Metadata filtering
- Collection management
- Async context manager support
"""
from typing import Any, Dict, List, Optional, Union, Callable
import uuid
import pickle
from pathlib import Path
import numpy as np
from navconfig.logging import logging
from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances

try:
    import faiss
    FAISS_AVAILABLE = True
except ImportError:
    FAISS_AVAILABLE = False
    faiss = None

from .abstract import AbstractStore
from .models import Document, SearchResult, DistanceStrategy


class FAISSStore(AbstractStore):
    """
    An in-memory FAISS vector store implementation, completely independent of Langchain.

    This store provides high-performance vector similarity search using FAISS indexes
    with support for multiple distance metrics and metadata filtering.

    Features:
    - Multiple FAISS index types (Flat, IVF, HNSW)
    - CPU-only execution
    - Cosine, L2, and Inner Product distance metrics
    - MMR (Maximal Marginal Relevance) search
    - Metadata filtering
    - Persistent storage via save/load
    """

    def __init__(
        self,
        collection_name: str = "default",
        id_column: str = 'id',
        embedding_column: str = 'embedding',
        document_column: str = 'document',
        text_column: str = 'text',
        metadata_column: str = 'metadata',
        embedding_model: Union[dict, str] = "sentence-transformers/all-mpnet-base-v2",
        embedding: Optional[Callable] = None,
        distance_strategy: DistanceStrategy = DistanceStrategy.COSINE,
        index_type: str = "Flat",  # Options: "Flat", "IVF", "HNSW"
        nlist: int = 100,  # For IVF indexes
        nprobe: int = 10,  # For IVF search
        m: int = 32,  # For HNSW
        ef_construction: int = 40,  # For HNSW
        ef_search: int = 16,  # For HNSW
        **kwargs
    ):
        """
        Initialize FAISSStore with the specified parameters.

        Args:
            collection_name: Name of the collection/index
            id_column: Name of the ID column
            embedding_column: Name of the embedding column
            document_column: Name of the document content column
            text_column: Name of the text column
            metadata_column: Name of the metadata column
            embedding_model: Embedding model configuration (dict or string)
            embedding: Custom embedding function
            distance_strategy: Distance metric to use (COSINE, EUCLIDEAN_DISTANCE, etc.)
            index_type: Type of FAISS index ("Flat", "IVF", "HNSW")
            nlist: Number of clusters for IVF indexes
            nprobe: Number of clusters to probe for IVF search
            m: Number of connections per layer for HNSW
            ef_construction: Size of dynamic candidate list for HNSW construction
            ef_search: Size of dynamic candidate list for HNSW search
        """
        if not FAISS_AVAILABLE:
            raise ImportError(
                "FAISS is not installed. Please install it with: pip install faiss-cpu"
            )

        # Store configuration
        self.collection_name = collection_name
        self._id_column: str = id_column
        self._embedding_column: str = embedding_column
        self._document_column: str = document_column
        self._text_column: str = text_column
        self._metadata_column: str = metadata_column

        # FAISS configuration
        self.index_type = index_type
        self.nlist = nlist
        self.nprobe = nprobe
        self.m = m
        self.ef_construction = ef_construction
        self.ef_search = ef_search

        # Distance strategy - normalize to enum
        if isinstance(distance_strategy, str):
            # Convert string to DistanceStrategy enum
            try:
                self.distance_strategy = DistanceStrategy[distance_strategy.upper()]
            except KeyError:
                self.logger.warning(
                    f"Unknown distance strategy '{distance_strategy}', using COSINE"
                )
                self.distance_strategy = DistanceStrategy.COSINE
        elif isinstance(distance_strategy, DistanceStrategy):
            self.distance_strategy = distance_strategy
        else:
            # Default to COSINE if invalid type
            self.logger.warning(
                f"Invalid distance_strategy type: {type(distance_strategy)}, using COSINE"
            )
            self.distance_strategy = DistanceStrategy.COSINE

        # Initialize parent class
        super().__init__(
            embedding_model=embedding_model,
            embedding=embedding,
            **kwargs
        )

        # Collections store: {collection_name: collection_data}
        self._collections: Dict[str, Dict[str, Any]] = {}

        # Initialize logger
        self.logger = logging.getLogger("FAISSStore")

        # Connection state
        self._connected: bool = False
        self._connection = None  # For compatibility with abstract interface

        # Initialize default collection
        if collection_name:
            self._initialize_collection(collection_name)

    def _initialize_collection(self, collection_name: str) -> None:
        """Initialize a new collection with empty data structures."""
        if collection_name not in self._collections:
            self._collections[collection_name] = {
                'index': None,  # FAISS index
                'documents': {},  # {id: document_content}
                'metadata': {},  # {id: metadata_dict}
                'embeddings': {},  # {id: embedding_vector}
                'id_to_idx': {},  # {id: faiss_index_position}
                'idx_to_id': {},  # {faiss_index_position: id}
                'dimension': None,
                'is_trained': False,
            }
            self.logger.info(f"Initialized collection: {collection_name}")

    def define_collection_table(
        self,
        collection_name: str,
        dimension: int = 384,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Define an in-memory collection table for saving vector + metadata information.

        This method is compatible with the PgVectorStore pattern but operates in-memory.

        Args:
            collection_name: Name of the collection
            dimension: Dimension of the embedding vectors
            **kwargs: Additional arguments (for compatibility)

        Returns:
            Dictionary representing the collection structure
        """
        if collection_name not in self._collections:
            self._initialize_collection(collection_name)

        collection = self._collections[collection_name]
        collection['dimension'] = dimension

        # Create FAISS index based on configuration
        index = self._create_faiss_index(dimension)
        collection['index'] = index

        self.logger.info(
            f"Defined collection table '{collection_name}' with dimension {dimension}"
        )

        return {
            'collection_name': collection_name,
            'dimension': dimension,
            'index_type': self.index_type,
            'distance_strategy': self.distance_strategy.value,
        }

    def _create_faiss_index(self, dimension: int) -> Any:
        """
        Create a FAISS index based on the configured index type and distance strategy.

        Args:
            dimension: Dimension of the embedding vectors

        Returns:
            FAISS index object
        """
        # Determine the metric type based on distance strategy
        if self.distance_strategy == DistanceStrategy.COSINE:
            # For cosine similarity, we'll normalize vectors and use inner product
            metric = faiss.METRIC_INNER_PRODUCT
        elif self.distance_strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
            metric = faiss.METRIC_L2
        elif self.distance_strategy in [DistanceStrategy.DOT_PRODUCT, DistanceStrategy.MAX_INNER_PRODUCT]:
            metric = faiss.METRIC_INNER_PRODUCT
        else:
            # Default to inner product
            metric = faiss.METRIC_INNER_PRODUCT

        # Create index based on type
        if self.index_type == "Flat":
            index = faiss.IndexFlatIP(dimension) if metric == faiss.METRIC_INNER_PRODUCT else faiss.IndexFlatL2(dimension)

        elif self.index_type == "IVF":
            # IVF (Inverted File Index) for faster search on large datasets
            quantizer = faiss.IndexFlatIP(dimension) if metric == faiss.METRIC_INNER_PRODUCT else faiss.IndexFlatL2(dimension)
            index = faiss.IndexIVFFlat(quantizer, dimension, self.nlist, metric)
            index.nprobe = self.nprobe

        elif self.index_type == "HNSW":
            # HNSW (Hierarchical Navigable Small World) for very fast search
            index = faiss.IndexHNSWFlat(dimension, self.m, metric)
            index.hnsw.efConstruction = self.ef_construction
            index.hnsw.efSearch = self.ef_search

        else:
            raise ValueError(f"Unsupported index type: {self.index_type}")

        self.logger.info(
            f"Created FAISS index: type={self.index_type}, "
            f"metric={metric}, dimension={dimension}, cpu_only=True"
        )

        return index

    async def connection(self) -> bool:
        """
        Establish connection (for compatibility with AbstractStore).

        Since FAISS is in-memory, this just marks the store as connected.

        Returns:
            True if successful
        """
        if not self._connected:
            self._connected = True
            self._connection = True  # Dummy connection for compatibility
            self.logger.info("FAISSStore connection established (in-memory)")
        return True

    async def disconnect(self) -> None:
        """
        Disconnect and cleanup resources.

        Clears all in-memory data.
        """
        if not self._connected:
            return
        # Clear indexes
        for _, collection in self._collections.items():
            if collection.get('index'):
                del collection['index']

        # Clear collections
        self._collections.clear()

        self._connected = False
        self._connection = None
        self.logger.info("FAISSStore disconnected and resources cleared")

    async def __aenter__(self):
        """Async context manager entry."""
        if not self._connected:
            await self.connection()
        self._context_depth += 1
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        """Async context manager exit."""
        # Free embedding resources
        if self._embed_:
            await self._free_resources()

        try:
            # Only disconnect if we're exiting the outermost context
            self._context_depth -= 1
            if self._context_depth <= 0:
                await self.disconnect()
                self._context_depth = 0
        except RuntimeError:
            pass

    async def prepare_embedding_table(
        self,
        collection: str = None,
        dimension: int = None,
        **kwargs
    ) -> bool:
        """
        Prepare the embedding table/collection for storing vectors.

        Args:
            collection: Collection name
            dimension: Embedding dimension
            **kwargs: Additional arguments

        Returns:
            True if successful
        """
        collection = collection or self.collection_name
        dimension = dimension or self.dimension

        if collection not in self._collections:
            self.define_collection_table(collection, dimension, **kwargs)
        else:
            # Update dimension if needed
            if dimension and self._collections[collection]['dimension'] != dimension:
                self._collections[collection]['dimension'] = dimension
                # Recreate index with new dimension
                self._collections[collection]['index'] = self._create_faiss_index(dimension)
                self.logger.info(f"Updated collection '{collection}' dimension to {dimension}")

        return True

    async def create_embedding_table(
        self,
        collection: str = None,
        dimension: int = None,
        **kwargs
    ) -> None:
        """
        Create an embedding table/collection (alias for prepare_embedding_table).

        Args:
            collection: Collection name
            dimension: Embedding dimension
            **kwargs: Additional arguments
        """
        await self.prepare_embedding_table(collection, dimension, **kwargs)

    async def create_collection(self, collection: str, **kwargs) -> None:
        """
        Create a new collection.

        Args:
            collection: Collection name
            **kwargs: Additional arguments (e.g., dimension)
        """
        dimension = kwargs.get('dimension', self.dimension)
        await self.create_embedding_table(collection, dimension, **kwargs)

    async def add_documents(
        self,
        documents: List[Document],
        collection: str = None,
        embedding_column: str = None,
        content_column: str = None,
        metadata_column: str = None,
        **kwargs
    ) -> None:
        """
        Add documents to the FAISS store.

        Args:
            documents: List of Document objects to add
            collection: Collection name (optional, uses default if not provided)
            embedding_column: Name of the embedding column (for compatibility)
            content_column: Name of the content column (for compatibility)
            metadata_column: Name of the metadata column (for compatibility)
            **kwargs: Additional arguments
        """
        if not self._connected:
            await self.connection()

        collection = collection or self.collection_name

        # Ensure collection exists
        if collection not in self._collections:
            self._initialize_collection(collection)

        collection_data = self._collections[collection]

        # Extract texts and metadata
        texts = [doc.page_content for doc in documents]
        metadatas = [doc.metadata for doc in documents]

        # Generate embeddings
        embeddings = await self._embed_.embed_documents(texts)

        # Convert to numpy array
        if isinstance(embeddings, list):
            embeddings = np.array(embeddings, dtype=np.float32)
        elif not isinstance(embeddings, np.ndarray):
            embeddings = np.array(embeddings, dtype=np.float32)

        # Ensure 2D array
        if embeddings.ndim == 1:
            embeddings = embeddings.reshape(1, -1)

        # Set dimension if not set
        if collection_data['dimension'] is None:
            collection_data['dimension'] = embeddings.shape[1]
            collection_data['index'] = self._create_faiss_index(embeddings.shape[1])

        # Normalize embeddings for cosine similarity
        if self.distance_strategy == DistanceStrategy.COSINE:
            faiss.normalize_L2(embeddings)

        # Train index if needed (for IVF)
        if self.index_type == "IVF" and not collection_data['is_trained']:
            if len(embeddings) >= self.nlist:
                collection_data['index'].train(embeddings)
                collection_data['is_trained'] = True
                self.logger.info(f"Trained IVF index for collection '{collection}'")
            else:
                self.logger.warning(
                    f"Not enough vectors to train IVF index "
                    f"(need {self.nlist}, got {len(embeddings)}). Using Flat index temporarily."
                )

        # Get current index size
        current_idx = collection_data['index'].ntotal

        # Add to FAISS index
        collection_data['index'].add(embeddings)

        # Store documents, metadata, and embeddings
        for i, (text, metadata, embedding) in enumerate(zip(texts, metadatas, embeddings)):
            doc_id = str(uuid.uuid4())
            idx = current_idx + i

            collection_data['documents'][doc_id] = text
            collection_data['metadata'][doc_id] = metadata or {}
            collection_data['embeddings'][doc_id] = embedding
            collection_data['id_to_idx'][doc_id] = idx
            collection_data['idx_to_id'][idx] = doc_id

        self.logger.info(
            f"✅ Successfully added {len(documents)} documents to collection '{collection}'"
        )

    def get_distance_strategy(
        self,
        query_embedding: np.ndarray,
        metric: str = None
    ) -> str:
        """
        Return the appropriate distance strategy based on the metric or configured strategy.

        Args:
            query_embedding: Query embedding vector (for compatibility)
            metric: Optional metric string ('COSINE', 'L2', 'IP', 'DOT')

        Returns:
            Distance strategy as string
        """
        strategy = metric or self.distance_strategy

        # Convert string metrics to DistanceStrategy enum if needed
        if isinstance(strategy, str):
            metric_mapping = {
                'COSINE': DistanceStrategy.COSINE,
                'L2': DistanceStrategy.EUCLIDEAN_DISTANCE,
                'EUCLIDEAN': DistanceStrategy.EUCLIDEAN_DISTANCE,
                'IP': DistanceStrategy.MAX_INNER_PRODUCT,
                'DOT': DistanceStrategy.DOT_PRODUCT,
                'DOT_PRODUCT': DistanceStrategy.DOT_PRODUCT,
                'MAX_INNER_PRODUCT': DistanceStrategy.MAX_INNER_PRODUCT
            }
            strategy = metric_mapping.get(strategy.upper(), DistanceStrategy.COSINE)

        return strategy

    async def similarity_search(
        self,
        query: str,
        collection: str = None,
        k: Optional[int] = None,
        limit: int = None,
        metadata_filters: Optional[Dict[str, Any]] = None,
        score_threshold: Optional[float] = None,
        metric: str = None,
        embedding_column: str = None,
        content_column: str = None,
        metadata_column: str = None,
        id_column: str = None,
        **kwargs
    ) -> List[SearchResult]:
        """
        Perform similarity search with optional threshold filtering.

        Args:
            query: The search query text
            collection: Collection name (optional, uses default if not provided)
            k: Number of results to return (alias for limit)
            limit: Maximum number of results to return
            metadata_filters: Dictionary of metadata filters to apply
            score_threshold: Minimum score threshold (results below threshold filtered out)
            metric: Distance metric to use ('COSINE', 'L2', 'IP')
            embedding_column: Name of the embedding column (for compatibility)
            content_column: Name of the content column (for compatibility)
            metadata_column: Name of the metadata column (for compatibility)
            id_column: Name of the ID column (for compatibility)
            **kwargs: Additional arguments

        Returns:
            List of SearchResult objects with content, metadata, score, and id
        """
        if not self._connected:
            await self.connection()

        collection = collection or self.collection_name

        if k and not limit:
            limit = k
        if not limit:
            limit = 10

        # Ensure collection exists
        if collection not in self._collections:
            self.logger.warning(f"Collection '{collection}' not found")
            return []

        collection_data = self._collections[collection]

        if collection_data['index'] is None or collection_data['index'].ntotal == 0:
            self.logger.warning(f"Collection '{collection}' is empty")
            return []

        # Generate query embedding
        query_embedding = await self._embed_.embed_query(query)

        # Convert to numpy array
        if isinstance(query_embedding, list):
            query_embedding = np.array(query_embedding, dtype=np.float32)
        elif not isinstance(query_embedding, np.ndarray):
            query_embedding = np.array(query_embedding, dtype=np.float32)

        # Ensure 2D array
        if query_embedding.ndim == 1:
            query_embedding = query_embedding.reshape(1, -1)

        # Normalize for cosine similarity
        if self.distance_strategy == DistanceStrategy.COSINE:
            faiss.normalize_L2(query_embedding)

        # Search FAISS index
        # For metadata filtering, we need to search more results and filter
        search_k = limit * 3 if metadata_filters else limit
        search_k = min(search_k, collection_data['index'].ntotal)

        distances, indices = collection_data['index'].search(query_embedding, search_k)

        # Convert distances to scores
        # FAISS returns distances, but we want similarity scores (higher is better)
        distances = distances[0]  # Get first row
        indices = indices[0]  # Get first row

        # Convert to scores based on metric
        if self.distance_strategy == DistanceStrategy.COSINE:
            # For cosine with normalized vectors, distance is actually similarity
            scores = distances
        elif self.distance_strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
            # Convert distance to similarity
            scores = 1.0 / (1.0 + distances)
        else:
            # For inner product, distance is already similarity-like
            scores = distances

        # Build results
        results = []
        for idx, score in zip(indices, scores):
            if idx == -1:  # FAISS returns -1 for empty slots
                continue

            # Get document ID
            doc_id = collection_data['idx_to_id'].get(idx)
            if doc_id is None:
                continue

            # Get document data
            content = collection_data['documents'].get(doc_id, "")
            metadata = collection_data['metadata'].get(doc_id, {})

            # Apply metadata filters
            if metadata_filters:
                match = True
                for key, value in metadata_filters.items():
                    if metadata.get(key) != value:
                        match = False
                        break
                if not match:
                    continue

            # Apply score threshold
            if score_threshold is not None:
                if self.distance_strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
                    # For L2 distance, lower is better, so invert the check
                    if distances[np.where(indices == idx)[0][0]] > score_threshold:
                        continue
                else:
                    # For cosine and inner product, higher is better
                    if score < score_threshold:
                        continue

            # Create search result
            result = SearchResult(
                id=doc_id,
                content=content,
                metadata=metadata,
                score=float(score)
            )
            results.append(result)

            # Stop if we have enough results
            if len(results) >= limit:
                break

        self.logger.debug(
            f"Similarity search in collection '{collection}': "
            f"found {len(results)} results (limit={limit})"
        )

        return results

    async def asearch(
        self,
        query: str,
        collection: Optional[str] = None,
        k: Optional[int] = None,
        limit: Optional[int] = None,
        metadata_filters: Optional[Dict[str, Any]] = None,
        score_threshold: Optional[float] = None,
        metric: Optional[str] = None,
        embedding_column: Optional[str] = None,
        content_column: Optional[str] = None,
        metadata_column: Optional[str] = None,
        id_column: Optional[str] = None,
        **kwargs,
    ) -> List[SearchResult]:
        """Async alias for :meth:`similarity_search` to match store interface expectations."""

        return await self.similarity_search(
            query=query,
            collection=collection,
            k=k,
            limit=limit,
            metadata_filters=metadata_filters,
            score_threshold=score_threshold,
            metric=metric,
            embedding_column=embedding_column,
            content_column=content_column,
            metadata_column=metadata_column,
            id_column=id_column,
            **kwargs,
        )

    async def mmr_search(
        self,
        query: str,
        collection: str = None,
        k: int = 4,
        fetch_k: Optional[int] = None,
        lambda_mult: float = 0.5,
        metadata_filters: Optional[Dict[str, Any]] = None,
        score_threshold: Optional[float] = None,
        metric: str = None,
        embedding_column: str = None,
        content_column: str = None,
        metadata_column: str = None,
        id_column: str = None,
        **kwargs
    ) -> List[SearchResult]:
        """
        Perform MMR (Maximal Marginal Relevance) search for diversity.

        MMR balances relevance with diversity by selecting documents that are:
        1. Relevant to the query
        2. Diverse from already selected documents

        Args:
            query: The search query text
            collection: Collection name (optional, uses default if not provided)
            k: Number of results to return
            fetch_k: Number of candidates to fetch (default: k*3)
            lambda_mult: Diversity parameter (0-1)
                - 1.0 = maximum relevance (no diversity)
                - 0.0 = maximum diversity (no relevance)
            metadata_filters: Dictionary of metadata filters to apply
            score_threshold: Minimum score threshold
            metric: Distance metric to use ('COSINE', 'L2', 'IP')
            embedding_column: Name of the embedding column (for compatibility)
            content_column: Name of the content column (for compatibility)
            metadata_column: Name of the metadata column (for compatibility)
            id_column: Name of the ID column (for compatibility)
            **kwargs: Additional arguments

        Returns:
            List of SearchResult objects ordered by MMR score
        """
        if not self._connected:
            await self.connection()

        collection = collection or self.collection_name

        # Default to fetching 3x more candidates than final results
        if fetch_k is None:
            fetch_k = max(k * 3, 20)

        # Step 1: Get initial candidates using similarity search
        candidates = await self.similarity_search(
            query=query,
            collection=collection,
            limit=fetch_k,
            metadata_filters=metadata_filters,
            score_threshold=score_threshold,
            metric=metric,
            embedding_column=embedding_column,
            content_column=content_column,
            metadata_column=metadata_column,
            id_column=id_column,
            **kwargs
        )

        if len(candidates) <= k:
            # If we have fewer candidates than requested results, return all
            return candidates

        # Step 2: Get embeddings for MMR computation
        collection_data = self._collections[collection]
        candidate_embeddings = {}
        for result in candidates:
            embedding = collection_data['embeddings'].get(result.id)
            if embedding is not None:
                candidate_embeddings[result.id] = embedding

        # Step 3: Get query embedding
        query_embedding = await self._embed_.embed_query(query)
        if isinstance(query_embedding, list):
            query_embedding = np.array(query_embedding, dtype=np.float32)
        elif not isinstance(query_embedding, np.ndarray):
            query_embedding = np.array(query_embedding, dtype=np.float32)

        # Step 4: Run MMR algorithm
        selected_results = self._mmr_algorithm(
            query_embedding=query_embedding,
            candidates=candidates,
            candidate_embeddings=candidate_embeddings,
            k=k,
            lambda_mult=lambda_mult,
            metric=metric or self.distance_strategy
        )

        self.logger.info(
            f"MMR search in collection '{collection}': "
            f"selected {len(selected_results)} results from {len(candidates)} candidates "
            f"(λ={lambda_mult})"
        )

        return selected_results

    def _mmr_algorithm(
        self,
        query_embedding: np.ndarray,
        candidates: List[SearchResult],
        candidate_embeddings: Dict[str, np.ndarray],
        k: int,
        lambda_mult: float,
        metric: Union[str, DistanceStrategy]
    ) -> List[SearchResult]:
        """
        Core MMR algorithm implementation (same as PgVectorStore).

        Args:
            query_embedding: Query embedding vector
            candidates: List of candidate SearchResult objects
            candidate_embeddings: Dictionary mapping doc ID to embedding vector
            k: Number of results to select
            lambda_mult: MMR diversity parameter (0-1)
            metric: Distance metric to use

        Returns:
            List of selected SearchResult objects ordered by MMR score
        """
        if len(candidates) <= k:
            return candidates

        # Convert query embedding to numpy array
        if not isinstance(query_embedding, np.ndarray):
            query_embedding = np.array(query_embedding, dtype=np.float32)

        # Prepare data structures
        selected_indices = []
        remaining_indices = list(range(len(candidates)))

        # Step 1: Select the most relevant document first
        query_similarities = []
        for candidate in candidates:
            doc_embedding = candidate_embeddings.get(candidate.id)
            if doc_embedding is not None:
                similarity = self._compute_similarity(query_embedding, doc_embedding, metric)
                query_similarities.append(similarity)
            else:
                # Fallback to distance score if embedding not available
                query_similarities.append(1.0 / (1.0 + candidate.score))

        # Select the most similar document first
        best_idx = np.argmax(query_similarities)
        selected_indices.append(best_idx)
        remaining_indices.remove(best_idx)

        # Step 2: Iteratively select remaining documents using MMR
        for _ in range(min(k - 1, len(remaining_indices))):
            mmr_scores = []

            for idx in remaining_indices:
                candidate = candidates[idx]
                doc_embedding = candidate_embeddings.get(candidate.id)

                if doc_embedding is None:
                    # Fallback scoring if embedding not available
                    mmr_score = lambda_mult * query_similarities[idx]
                    mmr_scores.append(mmr_score)
                    continue

                # Relevance: similarity to query
                relevance = query_similarities[idx]

                # Diversity: maximum similarity to already selected documents
                max_similarity_to_selected = 0.0
                for selected_idx in selected_indices:
                    selected_candidate = candidates[selected_idx]
                    selected_embedding = candidate_embeddings.get(selected_candidate.id)

                    if selected_embedding is not None:
                        similarity = self._compute_similarity(doc_embedding, selected_embedding, metric)
                        max_similarity_to_selected = max(max_similarity_to_selected, similarity)

                # MMR formula: λ * relevance - (1-λ) * max_similarity_to_selected
                mmr_score = (
                    lambda_mult * relevance -
                    (1.0 - lambda_mult) * max_similarity_to_selected
                )
                mmr_scores.append(mmr_score)

            # Select document with highest MMR score
            if mmr_scores:
                best_remaining_idx = np.argmax(mmr_scores)
                best_idx = remaining_indices[best_remaining_idx]
                selected_indices.append(best_idx)
                remaining_indices.remove(best_idx)

        # Step 3: Return selected results with MMR scores in metadata
        selected_results = []
        for i, idx in enumerate(selected_indices):
            result = candidates[idx]
            # Add MMR ranking to metadata
            enhanced_metadata = dict(result.metadata)
            enhanced_metadata['mmr_rank'] = i + 1
            enhanced_metadata['mmr_lambda'] = lambda_mult
            enhanced_metadata['original_rank'] = idx + 1

            enhanced_result = SearchResult(
                id=result.id,
                content=result.content,
                metadata=enhanced_metadata,
                score=result.score
            )
            selected_results.append(enhanced_result)

        return selected_results

    def _compute_similarity(
        self,
        embedding1: np.ndarray,
        embedding2: np.ndarray,
        metric: Union[str, DistanceStrategy]
    ) -> float:
        """
        Compute similarity between two embeddings based on the specified metric.

        Args:
            embedding1: First embedding vector (numpy array or list)
            embedding2: Second embedding vector (numpy array or list)
            metric: Distance metric ('COSINE', 'L2', 'IP', etc.)

        Returns:
            Similarity score (higher = more similar)
        """
        # Convert to numpy arrays if needed
        if isinstance(embedding1, list):
            embedding1 = np.array(embedding1, dtype=np.float32)
        if isinstance(embedding2, list):
            embedding2 = np.array(embedding2, dtype=np.float32)

        # Ensure embeddings are numpy arrays
        if not isinstance(embedding1, np.ndarray):
            embedding1 = np.array(embedding1, dtype=np.float32)
        if not isinstance(embedding2, np.ndarray):
            embedding2 = np.array(embedding2, dtype=np.float32)

        # Ensure embeddings are 2D arrays for sklearn
        emb1 = embedding1.reshape(1, -1)
        emb2 = embedding2.reshape(1, -1)

        # Convert string metrics to DistanceStrategy enum if needed
        if isinstance(metric, str):
            metric_mapping = {
                'COSINE': DistanceStrategy.COSINE,
                'L2': DistanceStrategy.EUCLIDEAN_DISTANCE,
                'EUCLIDEAN': DistanceStrategy.EUCLIDEAN_DISTANCE,
                'IP': DistanceStrategy.MAX_INNER_PRODUCT,
                'DOT': DistanceStrategy.DOT_PRODUCT,
                'DOT_PRODUCT': DistanceStrategy.DOT_PRODUCT,
                'MAX_INNER_PRODUCT': DistanceStrategy.MAX_INNER_PRODUCT
            }
            strategy = metric_mapping.get(metric.upper(), DistanceStrategy.COSINE)
        else:
            strategy = metric

        if strategy == DistanceStrategy.COSINE:
            # Cosine similarity (returns similarity, not distance)
            similarity = cosine_similarity(emb1, emb2)[0, 0]

        elif strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
            # Convert Euclidean distance to similarity
            distance = euclidean_distances(emb1, emb2)[0, 0]
            similarity = 1.0 / (1.0 + distance)

        elif strategy in [DistanceStrategy.MAX_INNER_PRODUCT, DistanceStrategy.DOT_PRODUCT]:
            # Dot product (inner product)
            similarity = np.dot(embedding1.flatten(), embedding2.flatten())

        else:
            # Default to cosine similarity
            similarity = cosine_similarity(emb1, emb2)[0, 0]
        return float(similarity)

    # Additional methods for compatibility

    def get_vector(self, metric_type: str = None, **kwargs):
        """
        Get the FAISS vector store (for compatibility).

        Args:
            metric_type: Distance metric type
            **kwargs: Additional arguments

        Returns:
            The FAISSStore instance itself
        """
        return self

    async def from_documents(
        self,
        documents: List[Document],
        collection: Union[str, None] = None,
        **kwargs
    ) -> 'FAISSStore':
        """
        Create Vector Store from Documents.

        Args:
            documents: List of Documents
            collection: Collection Name
            **kwargs: Additional Arguments

        Returns:
            The FAISSStore instance
        """
        await self.add_documents(documents, collection=collection, **kwargs)
        return self

    # Persistence methods

    def save(self, file_path: Union[str, Path]) -> None:
        """
        Save the FAISS store to disk.

        Args:
            file_path: Path to save the store
        """
        file_path = Path(file_path)
        file_path.parent.mkdir(parents=True, exist_ok=True)

        # Prepare data for saving
        save_data = {
            'collections': {},
            'config': {
                'collection_name': self.collection_name,
                'distance_strategy': self.distance_strategy.value,
                'index_type': self.index_type,
                'dimension': self.dimension,
            }
        }

        # Save each collection
        for coll_name, coll_data in self._collections.items():
            # Save FAISS index
            index_path = file_path.parent / f"{file_path.stem}_{coll_name}.index"
            if coll_data['index'] is not None:
                faiss.write_index(coll_data['index'], str(index_path))

            # Save metadata and documents
            save_data['collections'][coll_name] = {
                'documents': coll_data['documents'],
                'metadata': coll_data['metadata'],
                'embeddings': {k: v.tolist() for k, v in coll_data['embeddings'].items()},
                'id_to_idx': coll_data['id_to_idx'],
                'idx_to_id': coll_data['idx_to_id'],
                'dimension': coll_data['dimension'],
                'is_trained': coll_data['is_trained'],
                'index_path': str(index_path),
            }

        # Save to pickle
        with open(file_path, 'wb') as f:
            pickle.dump(save_data, f)

        self.logger.info(f"Saved FAISS store to {file_path}")

    def load(self, file_path: Union[str, Path]) -> None:
        """
        Load the FAISS store from disk.

        Args:
            file_path: Path to load the store from
        """
        file_path = Path(file_path)

        with open(file_path, 'rb') as f:
            save_data = pickle.load(f)

        # Restore config
        config = save_data.get('config', {})
        self.collection_name = config.get('collection_name', self.collection_name)

        # Restore collections
        for coll_name, coll_data in save_data['collections'].items():
            self._initialize_collection(coll_name)

            # Load FAISS index
            index_path = coll_data['index_path']
            if Path(index_path).exists():
                index = faiss.read_index(index_path)

                self._collections[coll_name]['index'] = index

            # Load metadata and documents
            self._collections[coll_name]['documents'] = coll_data['documents']
            self._collections[coll_name]['metadata'] = coll_data['metadata']
            self._collections[coll_name]['embeddings'] = {
                k: np.array(v, dtype=np.float32)
                for k, v in coll_data['embeddings'].items()
            }
            self._collections[coll_name]['id_to_idx'] = coll_data['id_to_idx']
            self._collections[coll_name]['idx_to_id'] = coll_data['idx_to_id']
            self._collections[coll_name]['dimension'] = coll_data['dimension']
            self._collections[coll_name]['is_trained'] = coll_data['is_trained']

        self.logger.info(f"Loaded FAISS store from {file_path}")

    def __str__(self) -> str:
        return f"FAISSStore(collection={self.collection_name}, index_type={self.index_type})"

    def __repr__(self) -> str:
        return (
            f"<FAISSStore(collection='{self.collection_name}', "
            f"index_type='{self.index_type}', "
            f"distance_strategy='{self.distance_strategy.value}', "
            "cpu_only=True)>"
        )

    async def delete_documents(
        self,
        document_ids: List[str],
        collection: str = None,
        **kwargs
    ) -> None:
        """
        Delete documents by their IDs from the FAISS store.

        Args:
            document_ids: List of document IDs to delete
            collection: Collection name (optional, uses default if not provided)
            **kwargs: Additional arguments
        """
        if not self._connected:
            await self.connection()

        collection = collection or self.collection_name

        # Ensure collection exists
        if collection not in self._collections:
            self.logger.warning(f"Collection '{collection}' not found")
            return

        collection_data = self._collections[collection]

        for doc_id in document_ids:
            idx = collection_data['id_to_idx'].get(doc_id)
            if idx is not None:
                # Remove from FAISS index
                # Note: FAISS does not support direct deletion; we mark as deleted
                # Here we simply ignore the vector in searches by removing mappings
                del collection_data['documents'][doc_id]
                del collection_data['metadata'][doc_id]
                del collection_data['embeddings'][doc_id]
                del collection_data['id_to_idx'][doc_id]
                del collection_data['idx_to_id'][idx]

        self.logger.info(
            f"✅ Successfully deleted {len(document_ids)} documents from collection '{collection}'"
        )

    async def delete_documents_by_filter(self, filter_func, collection: str = None, **kwargs) -> None:
        """
        Delete documents that match a filter function from the FAISS store.

        Args:
            filter_func: A function that takes metadata and returns True if the document should be deleted
            collection: Collection name (optional, uses default if not provided)
            **kwargs: Additional arguments
        """
        if not self._connected:
            await self.connection()

        collection = collection or self.collection_name

        # Ensure collection exists
        if collection not in self._collections:
            self.logger.warning(f"Collection '{collection}' not found")
            return

        collection_data = self._collections[collection]

        to_delete_ids = []
        for doc_id, metadata in collection_data['metadata'].items():
            if filter_func(metadata):
                to_delete_ids.append(doc_id)

        await self.delete_documents(to_delete_ids, collection=collection, **kwargs)
