"""Dual vector store implementation for ingesting documents into both sparse and dense stores"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/01_ingest.stores.dual.ipynb.

# %% auto 0
__all__ = ['DualStore', 'ElasticsearchStore']

# %% ../../../nbs/01_ingest.stores.dual.ipynb 3
import os
from typing import List, Optional, Callable, Dict, Sequence, Union
from tqdm import tqdm
from langchain_core.documents import Document

from .base import VectorStore
from .dense import DenseStore
from .sparse import SparseStore, ElasticsearchSparseStore, ELASTICSEARCH_INSTALLED
from ..helpers import doc_from_dict

class DualStore(VectorStore):
    def __init__(
        self,
        dense_kind:str='chroma',
        dense_persist_location: Optional[str] = None,
        sparse_kind:str='whoosh',
        sparse_persist_location: Optional[str] = None,
        **kwargs
    ):
        """
        Initialize a dual vector store that manages both dense and sparse stores.
        
        **Args**:
        
          - *dense_persist_location*: Path to dense vector database (created if it doesn't exist).
          - *sparse_persist_location*: Path to sparse vector database (created if it doesn't exist).
          - *embedding_model_name*: name of sentence-transformers model
          - *embedding_model_kwargs*: arguments to embedding model (e.g., `{device':'cpu'}`). If None, GPU used if available.
          - *embedding_encode_kwargs*: arguments to encode method of embedding model (e.g., `{'normalize_embeddings': False}`).
        """
        self.init_embedding_model(**kwargs)  # stored in self.embeddings
        
        # Extract the embedding parameters that were actually used
        embedding_model_name = kwargs.get('embedding_model_name', "sentence-transformers/all-MiniLM-L6-v2")
        embedding_model_kwargs = kwargs.get('embedding_model_kwargs')
        embedding_encode_kwargs = kwargs.get('embedding_encode_kwargs', {"normalize_embeddings": False})
        
        # Initialize both stores with consistent embedding parameters
        self.dense_store = DenseStore.create(
            kind=dense_kind,
            persist_location=dense_persist_location,
            embedding_model_name=embedding_model_name,
            embedding_model_kwargs=embedding_model_kwargs,
            embedding_encode_kwargs=embedding_encode_kwargs
        )
        self.sparse_store = SparseStore.create(
            kind=sparse_kind,
            persist_location=sparse_persist_location,
            embedding_model_name=embedding_model_name,
            embedding_model_kwargs=embedding_model_kwargs,
            embedding_encode_kwargs=embedding_encode_kwargs
        )
        
        # For compatibility with the VectorStore interface
        self.persist_location = dense_persist_location

    @classmethod
    def create(cls, dense_kind='chroma', sparse_kind='whoosh', persist_location=None, **kwargs):
        """
        Factory method to construct a DualStore instance.
        
        Args:
            dense_kind: type of dense store ('chroma', 'elasticsearch')
            sparse_kind: type of sparse store ('whoosh', 'elasticsearch')
            persist_location: base directory for stores or Elasticsearch URL
            **kwargs: additional arguments passed to store initialization
            
        For traditional dual stores (different dense_kind and sparse_kind):
            dense_persist_location: directory for dense store  
            sparse_persist_location: directory for sparse store
            
        For unified Elasticsearch dual store (both kinds are 'elasticsearch'):
            index_name: name of Elasticsearch index
            dense_vector_field: field name for dense vectors
            basic_auth: authentication credentials
            verify_certs: SSL verification
            ca_certs: CA certificate path
            timeout: connection timeout
            
        Returns:
            DualStore instance
        """
        # If both dense and sparse are elasticsearch, use unified ElasticsearchDualStore
        if dense_kind == 'elasticsearch' and sparse_kind == 'elasticsearch':
            if not ELASTICSEARCH_INSTALLED:
                raise ImportError('Please install elasticsearch packages: pip install onprem[elasticsearch]')
            return ElasticsearchStore(persist_location=persist_location, **kwargs)
        
        # Otherwise, use traditional dual store approach
        return cls(
            dense_kind=dense_kind,
            sparse_kind=sparse_kind,
            dense_persist_location=persist_location,
            sparse_persist_location=persist_location,
            **kwargs
        )


    # get_dense_db() and get_sparse_db() methods removed - use store methods instead


    #------------------------------
    # overrides of abstract methods
    # -----------------------------
  
    
    def exists(self):
        """
        Returns True if either store exists.
        """
        return self.dense_store.exists() or self.sparse_store.exists()
    
    def add_documents(self, documents: Sequence[Document], batch_size: int = 1000, **kwargs):
        """
        Add documents to both dense and sparse stores.
        If both stores use the same persist_location, only add once.
        """
        if not documents:
            return
        
        # Check if both stores are using the same persist_location
        if self._stores_share_same_location():
            # Use the dense store since it typically includes both text and vectors
            self.dense_store.add_documents(documents, batch_size=batch_size, **kwargs)
        else:
            # Traditional dual store approach - add to both stores
            self.dense_store.add_documents(documents, batch_size=batch_size, **kwargs)
            self.sparse_store.add_documents(documents, **kwargs)
    
    def _stores_share_same_location(self):
        """
        Check if the dense and sparse stores are using the same persist_location.
        This indicates they're using the same underlying storage/index.
        """
        dense_location = getattr(self.dense_store, 'persist_location', None)
        sparse_location = getattr(self.sparse_store, 'persist_location', None)
        
        return (dense_location is not None and 
                sparse_location is not None and 
                dense_location == sparse_location)

   
    def remove_document(self, id_to_delete):
        """
        Remove a document from both stores.
        """
        self.dense_store.remove_document(id_to_delete)
        self.sparse_store.remove_document(id_to_delete)

    def remove_source(self, source:str):
        """
        Remove a document by source from both stores.

        The `source` can either be the full path to a document
        or a parent folder.  Returns the number of records deleted.
        """
        try:
            num_deleted_1 = self.dense_store.remove_source(source)
            num_deleted_2 = self.sparse_store.remove_source(source)
            return num_deleted_1
        except ValueError as e:
            # Re-raise with more context if source field is not configured
            if "no source field configured" in str(e):
                raise ValueError("Cannot remove by source: neither dense nor sparse store has a source field configured. Set source_field parameter when creating the store.")
            raise
    

    def update_documents(self, doc_dicts: dict, **kwargs):
        """
        Update documents in both stores.
        """
        self.dense_store.update_documents(doc_dicts, **kwargs)
        self.sparse_store.update_documents(doc_dicts, **kwargs)
    
    def get_all_docs(self):
        """
        Get all documents from the dense store.
        For simplicity, we only return documents from one store since they should be the same.
        """
        return self.dense_store.get_all_docs()
    
    def get_doc(self, id):
        """
        Get a document by ID from the dense store.
        """
        return self.dense_store.get_doc(id)
    
    def get_size(self):
        """
        Get the size of the dense store.
        """
        return self.dense_store.get_size()
    
    def erase(self, confirm=True):
        """
        Erase both stores.
        """
        dense_erased = self.dense_store.erase(confirm=confirm)
        sparse_erased = self.sparse_store.erase(confirm=False)  # Second confirmation not needed
        return dense_erased and sparse_erased
    
    def search(self, q: str, **kwargs):
        """
        Query using the sparse store.
        """
        return self.sparse_store.search(q, **kwargs)
    
    def semantic_search(self, query: str, **kwargs):
        """
        Perform semantic search using the dense store.
        """
        return self.dense_store.semantic_search(query, **kwargs)
    
    def hybrid_search(self, query: str, limit: int = 10, weights: Union[List[float], float] = 0.5, **kwargs):
        """
        Perform hybrid search combining dense and sparse results.
        
        **Args**:
        - *query*: Search query string
        - *limit*: Maximum number of results to return
        - *weights*: Weights for combining dense and sparse scores. 
                    If float, represents dense weight (sparse = 1 - dense).
                    If list, [dense_weight, sparse_weight]
        - *kwargs*: Additional arguments passed to individual search methods
        
        **Returns**:
        Dictionary with 'hits' and 'total_hits' keys
        """
        # Create weights array if single number passed
        if isinstance(weights, (int, float)):
            weights = [weights, 1 - weights]
        
        # Get expanded results from both stores
        search_limit = limit * 10  # Get more candidates for better fusion
        
        # Get dense results
        dense_results = self.dense_store.search(query, limit=search_limit, **kwargs)
        dense_hits = dense_results.get('hits', [])
        
        # Get sparse results  
        sparse_results = self.sparse_store.search(query, limit=search_limit, return_dict=True, **kwargs)
        sparse_hits = sparse_results.get('hits', [])
        
        # Combine scores using hybrid approach
        uids = {}
        
        # Process dense results (similarity-based scores)
        if weights[0] > 0:
            for dense_doc in dense_hits:
                uid = dense_doc.get('id')
                if uid:
                    score = dense_doc.get('score', 0.0)
                    uids[uid] = score * weights[0]
        
        # Process sparse results (use RRF since sparse doesn't have normalized scores)
        if weights[1] > 0:
            for rank, sparse_doc in enumerate(sparse_hits):
                uid = sparse_doc.get('id')
                if uid:
                    # Use Reciprocal Rank Fusion (RRF) for sparse results
                    rrf_score = 1.0 / (rank + 1)
                    if uid in uids:
                        uids[uid] += rrf_score * weights[1]
                    else:
                        uids[uid] = rrf_score * weights[1]
        
        # Sort by combined score and limit results
        sorted_results = sorted(uids.items(), key=lambda x: x[1], reverse=True)[:limit]
        
        # Convert back to result dictionaries
        results = []
        # Create a lookup dict for efficient document retrieval
        doc_lookup = {}
        for doc in dense_hits + sparse_hits:
            doc_lookup[doc.get('id')] = doc
        
        for uid, combined_score in sorted_results:
            if uid in doc_lookup:
                doc_dict = doc_lookup[uid].copy()
                doc_dict['score'] = combined_score  # Update with combined score
                results.append(doc_dict)
        
        return [doc_from_dict(r) for r in results] 


class ElasticsearchStore(DualStore):
    """
    A unified Elasticsearch-based dual store that supports both dense vector searches
    and sparse text searches in a single index. Uses composition to manage both stores.
    """
    
    def __init__(self,
                dense_vector_field: str = 'dense_vector',
                **kwargs):
        """
        Initializes unified Elasticsearch dual store with both dense and sparse capabilities.
        
        **Args:**
        - *dense_vector_field*: field name for dense vectors (default: 'dense_vector')
        - All other args are passed to both stores (persist_location, index_name, basic_auth, etc.)
        """
        # Initialize embedding model (required by DualStore interface)
        self.init_embedding_model(**kwargs)
        
        # Create both stores pointing to the same Elasticsearch index
        from .dense import ElasticsearchDenseStore
        self.dense_store = ElasticsearchDenseStore(
            dense_vector_field=dense_vector_field, **kwargs)
        self.sparse_store = ElasticsearchSparseStore(**kwargs)
        
        # For compatibility with the VectorStore interface
        self.persist_location = kwargs.get('persist_location')

    # get_dense_db() and get_sparse_db() methods removed - use store methods instead

    # hybrid_search method is inherited from DualStore
