"""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']

# %% ../../../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
from ..helpers import doc_from_dict

class DualStore(VectorStore):
    def __init__(
        self,
        dense_persist_directory: Optional[str] = None,
        sparse_persist_directory: Optional[str] = None,
        **kwargs
    ):
        """
        Initialize a dual vector store that manages both dense and sparse stores.
        
        **Args**:
        
          - *dense_persist_directory*: Path to dense vector database (created if it doesn't exist).
          - *sparse_persist_directory*: 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
        
        # Initialize both stores
        self.dense_store = DenseStore(
            persist_directory=dense_persist_directory,
            embedding_model_name=kwargs.get('embedding_model_name'),
            embedding_model_kwargs=kwargs.get('embedding_model_kwargs'),
            embedding_encode_kwargs=kwargs.get('embedding_encode_kwargs')
        )
        self.sparse_store = SparseStore(
            persist_directory=sparse_persist_directory,
            embedding_model_name=kwargs.get('embedding_model_name'),
            embedding_model_kwargs=kwargs.get('embedding_model_kwargs'),
            embedding_encode_kwargs=kwargs.get('embedding_encode_kwargs')
        )
        
        # For compatibility with the VectorStore interface
        self.persist_directory = dense_persist_directory
        
    def get_db(self):
        """
        Returns the dense store's database instance.
        For consistency with the VectorStore interface.
        """
        return self.dense_store.get_db()
    
    def get_dense_db(self):
        """
        Returns the dense store's database instance.
        """
        return self.dense_store.get_db()
    
    def get_sparse_db(self):
        """
        Returns the sparse store's database instance.
        """
        return self.sparse_store.get_db()
    
    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 not documents:
            return
        # Add to dense store
        self.dense_store.add_documents(documents, batch_size=batch_size, **kwargs)
        
        # Add to sparse store
        self.sparse_store.add_documents(documents, **kwargs)
    
    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.
        """
        num_deleted_1 = self.dense_store.remove_source(source)
        num_deleted_2 = self.sparse_store.remove_source(source)
        return num_deleted_1


    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 query(self, q: str, **kwargs):
        """
        Query using the sparse store.
        """
        return self.sparse_store.query(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 keyword_search(self, query: str, **kwargs):
        """
        Perform keyword search using the sparse store.
        """
        return self.sparse_store.query(query, **kwargs)
        
    def optimize_for_search(self, ef: int = 200):
        """
        Optimize the dense store's HNSW index parameters for search.
        This helps fix "ef or M is too small" errors when filtering by folders.
        
        Args:
            ef: The ef parameter value for HNSW search (default: 200, higher = more accurate but slower)
            
        Returns:
            True if optimization succeeded, False otherwise
        """
        if hasattr(self.dense_store, 'optimize_for_search'):
            return self.dense_store.optimize_for_search(ef=ef)
        return False
