"""Data store for dense vector represenations of documents"""

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

# %% auto 0
__all__ = ['COLLECTION_NAME', 'DenseStore', 'ChromaStore', 'ElasticsearchDenseStore']

# %% ../../../nbs/01_ingest.stores.dense.ipynb 3
from .base import VectorStore


class DenseStore(VectorStore):
    """
    A factory for built-in DenseStore instances.
    """
    def __init__(self, **kwargs):
        if type(self) is DenseStore:
            raise TypeError("Use the DenseStore.create() method instead of instantiating DenseStore directly.")

    @classmethod
    def create(cls, persist_location=None, kind=None, **kwargs) -> 'DenseStore':
        """
        Factory method to construct a `DenseStore` instance. 
        
        Extra kwargs passed to object instantiation.
        
        Args:
            persist_location: where the vector database is stored
            kind: one of {chroma, elasticsearch}

        Returns:
            DenseStore instance
        """
        
        kind = 'chroma' if not kind else kind
        
        if kind == 'chroma':
            return ChromaStore(persist_location=persist_location, **kwargs)
        elif kind == 'elasticsearch':
            return ElasticsearchDenseStore(persist_location=persist_location, **kwargs)
        else:
            raise ValueError(f"Unknown DenseStore type: {kind}. Supported types: 'chroma', 'elasticsearch'")

# %% ../../../nbs/01_ingest.stores.dense.ipynb 5
import os
import os.path
from typing import List, Optional, Callable, Dict, Sequence
from tqdm import tqdm
from langchain_core.documents import Document

from ..helpers import doc_from_dict, dict_from_doc
from ...utils import get_datadir, DEFAULT_DB
from ..base import batchify_chunks, process_folder, does_vectorstore_exist
from ..base import CHROMA_MAX
try:
    from langchain_chroma import Chroma
    import chromadb
    from chromadb.config import Settings
    CHROMA_INSTALLED = True
except ImportError:
    CHROMA_INSTALLED = False


os.environ["TOKENIZERS_PARALLELISM"] = "0"
COLLECTION_NAME = "onprem_chroma"


class ChromaStore(DenseStore):
    """
    A dense vector store based on Chroma.   
    """
    def __init__(
        self,
        persist_location: Optional[str] = None,
        **kwargs
    ):
        """
        Ingests all documents in `source_folder` (previously-ingested documents are ignored)

        **Args**:

          - *persist_location*: Path to vector database (created if it doesn't exist).
                                 Default is `onprem_data/vectordb` in user's home directory.
          - *embedding_model*: 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}`).


        **Returns**: `None`
        """
        if not CHROMA_INSTALLED:
            raise ImportError('Please install chroma packages: pip install onprem[chroma]')

        from langchain_chroma import Chroma
        import chromadb
        from chromadb.config import Settings

        self.persist_location = persist_location or os.path.join(
            get_datadir(), DEFAULT_DB
        )
        self.init_embedding_model(**kwargs) # stored in self.embeddings

        self.chroma_settings = Settings(
            persist_directory=self.persist_location, anonymized_telemetry=False
        )
        self.chroma_client = chromadb.PersistentClient(
            settings=self.chroma_settings, path=self.persist_location
        )
        return


    def _convert_to_dict(self, raw_results):
        """
        Convert raw results to dictionary
        """
        ids = raw_results['ids']
        texts = raw_results['documents']
        metadatas = raw_results['metadatas']
        results = []
        for i, m in enumerate(metadatas):
            m['page_content'] = texts[i]
            m['id'] = ids[i]
            results.append(m)
        return results

    def get_db(self):
        """
        Returns an instance to the `langchain_chroma.Chroma` instance
        """
        # Create ChromaDB settings
        db = Chroma(
            persist_directory=self.persist_location,
            embedding_function=self.embeddings,
            client_settings=self.chroma_settings,
            client=self.chroma_client,
            collection_metadata={"hnsw:space": "cosine"},
            collection_name=COLLECTION_NAME,
        )
        return db if does_vectorstore_exist(db) else None
    
    # get_db() method removed - use store methods instead
    
    #------------------------------
    # overrides of abstract methods
    # -----------------------------
    

    def exists(self):
        return self.get_db() is not None


    def add_documents(self, documents, batch_size:int=CHROMA_MAX):
        """
        Stores instances of `langchain_core.documents.base.Document` in vectordb
        """
        if not documents:
            return
        db = self.get_db()
        if db:
            print("Creating embeddings. May take some minutes...")
            chunk_batches, total_chunks = batchify_chunks(documents, batch_size=batch_size)
            for lst in tqdm(chunk_batches, total=total_chunks):
                db.add_documents(lst)
        else:
            chunk_batches, total_chunks = batchify_chunks(documents, batch_size)
            print("Creating embeddings. May take some minutes...")
            db = None

            for lst in tqdm(chunk_batches, total=total_chunks):
                if not db:
                    db = Chroma.from_documents(
                        lst,
                        self.embeddings,
                        persist_directory=self.persist_location,
                        client_settings=self.chroma_settings,
                        client=self.chroma_client,
                        collection_metadata={"hnsw:space": "cosine"},
                        collection_name=COLLECTION_NAME,
                    )
                else:
                    db.add_documents(lst)
        return


    def remove_document(self, id_to_delete):
        """
        Remove a single document with ID, `id_to_delete`.
        """
        if not self.exists(): return
        id_to_delete = [id_to_delete] if not isinstance(id_to_delete, list) else id_to_delete
        self.get_db().delete(ids=id_to_delete)
        return

    def remove_source(self, source:str):
        """
        Deletes all documents in a Chroma collection whose `source` metadata field starts with the given prefix.
        The `source` argument can either be a full path to a document or a prefix (e.g., parent folder).

        **Args:**
        - *source*: The source value or prefix

        **Returns:**
        - The number of documents deleted
        """
        db = self.get_db()

        # Only request metadata; ids are returned automatically
        results = db.get(include=["metadatas"])

        to_delete = []
        for doc_id, metadata in zip(results["ids"], results["metadatas"]):
            if metadata and "source" in metadata:
                if metadata["source"].startswith(source):
                    to_delete.append(doc_id)

        if to_delete:
            db.delete(ids=to_delete)
            return len(to_delete)
        else:
            return 0
    
    def update_documents(self,
                         doc_dicts:dict, # dictionary with keys 'page_content', 'source', 'id', etc.
                         **kwargs):

        """
        Update a set of documents (doc in index with same ID will be over-written)
        """
        self.check()
        db = self.get_db()
        docs = [doc_from_dict(d) for d in doc_dicts]
        ids = [d['id'] for d in doc_dicts]
        return db.update_documents(ids, docs)


   
    def get_all_docs(self):
        """
        Returns all docs
        """
        if not self.exists(): return []

        raw_results =  self.get_db().get()
        return self._convert_to_dict(raw_results)


    def get_doc(self, id):
        """
        Retrieve a record by ID
        """
        if not self.exists(): return None
        raw_results = self.get_db().get(ids=[id])
        return self._convert_to_dict(raw_results)[0] if len(raw_results['ids']) > 0 else None

    
    def get_size(self):
        """
        Get total number of records
        """
        if not self.exists(): return 0
        return len(self.get_db().get()['documents'])

    
    def erase(self, confirm=True):
        """
        Resets collection and removes and stored documents
        """
        if not self.exists(): return True
        shall = True
        if confirm:
            msg = (
                f"You are about to remove all documents from the vector database."
                + f"(Original documents on file system will remain.) Are you sure?"
            )
            shall = input("%s (Y/n) " % msg) == "Y"
        if shall:
            self.get_db().reset_collection()
            return True
        return False


    def search(self,
              query:str, # query string
              limit:int = 4, # max number of results to return
              filters:Optional[Dict[str, str]] = None, # filter sources by metadata values using Chroma metadata syntax (e.g., {'table':True})
              where_document:Optional[Dict[str, str]] = None, # filter sources by document content in Chroma syntax (e.g., {"$contains": "Canada"}),
              return_dict:bool=True, # If True, return results as dictionaries. Otherwise, return LangChain Document objects.
              **kwargs):
        """
        Perform a semantic search of the vector DB. Returns results as dictionary by default.
        """
        if not self.exists(): return []
        db = self.get_db()
        
        # Convert standard dict filters to ChromaDB $and format when multiple keys present
        chroma_filters = filters
        if filters and len(filters) > 1:
            chroma_filters = {
                "$and": [{k: v} for k, v in filters.items()]
            }
        
        results = db.similarity_search_with_score(query, 
                                                  filter=chroma_filters,
                                                  where_document=where_document,
                                                  k = limit, **kwargs)
        if not results: return []
        docs, scores = zip(*results)
        for doc, score in zip(docs, scores):
            simscore = 1 - score
            doc.metadata["score"] = 1-score
        return {'hits' : [dict_from_doc(d) for d in docs], 'total_hits': len(docs)} if return_dict else docs      


    def semantic_search(self, *args, **kwargs):
        """
        Perform a semantic search of the vector DB. Returns results as LangChain Document objects.
        """
        return self.search(*args, return_dict=False, **kwargs)
    

class ElasticsearchDenseStore(DenseStore):
    """
    Elasticsearch store with dense vector search capabilities.
    Extends DenseStore to provide Elasticsearch-based dense vector storage.
    """
    
    def __init__(self,
                dense_vector_field: str = 'dense_vector',
                **kwargs):
        """
        Initialize Elasticsearch dense store.
        
        **Args:**
        - *dense_vector_field*: field name for dense vectors (default: 'dense_vector')
        - All other args are passed to ElasticsearchStore for the underlying implementation
        """
        from .sparse import ElasticsearchSparseStore
        
        self.dense_vector_field = dense_vector_field
        
        # Create underlying ElasticsearchSparseStore instance
        self.es_store = ElasticsearchSparseStore(**kwargs)
        
        # Set persist_location to match the underlying store
        self.persist_location = self.es_store.persist_location
        
        # Initialize embedding model
        self.init_embedding_model(**kwargs)
        
        # Set up index with vector mappings
        self._setup_vector_index()
    
    def _setup_vector_index(self):
        """Ensure the index has vector field mappings"""
        # Check if index exists
        if not self.es_store.es.indices.exists(index=self.es_store.index_name):
            # Create index with vector mappings using custom field names from es_store
            properties = {
                # Essential fields for core functionality using custom field names and analyzer
                self.es_store.content_field: {"type": "text", "analyzer": self.es_store.content_analyzer},
                self.es_store.id_field: {"type": "keyword"},
                
                # Dense vector field for semantic search
                self.dense_vector_field: {
                    "type": "dense_vector",
                    "index": True,
                    "similarity": "cosine"
                }
            }
            
            # Add source field only if specified
            if self.es_store.source_field:
                properties[self.es_store.source_field] = {"type": "keyword"}
                properties[f"{self.es_store.source_field}_search"] = {"type": "text", "analyzer": self.es_store.content_analyzer}
            
            mapping = {
                "mappings": {
                    "properties": properties
                }
            }
            
            self.es_store.es.indices.create(index=self.es_store.index_name, body=mapping)
        else:
            # Index exists, update mapping to include vector field if needed
            try:
                mapping = {
                    "properties": {
                        self.dense_vector_field: {
                            "type": "dense_vector",
                            "index": True,
                            "similarity": "cosine"
                        }
                    }
                }
                self.es_store.es.indices.put_mapping(index=self.es_store.index_name, body=mapping)
            except Exception:
                # Mapping update failed, but that's okay if field already exists
                pass
    
    def doc2dict(self, doc: Document, include_vector: bool = True):
        """Convert LangChain Document to expected format with vector embedding"""
        # Get the standard dict from ElasticsearchSparseStore
        d = self.es_store.doc2dict(doc)
        
        # Add dense vector embedding if requested
        if include_vector and hasattr(self, 'embeddings'):
            try:
                # Generate embedding for the document text
                embedding = self.embeddings.embed_documents([doc.page_content])[0]
                d[self.dense_vector_field] = embedding
            except Exception as e:
                # If embedding fails, continue without it
                pass
                
        return d
    
    # Override abstract methods from DenseStore
    def exists(self):
        """Returns True if the Elasticsearch index exists and has documents"""
        return self.es_store.exists()
    
    def add_documents(self, documents, batch_size: int = 1000, **kwargs):
        """Add documents with vector embeddings to Elasticsearch"""
        if not documents:
            return
        
        # Convert documents to dict format with vectors
        actions = []
        for doc in documents:
            d = self.doc2dict(doc, include_vector=True)
            action = {
                "_index": self.es_store.index_name,
                "_id": d[self.es_store.id_field],
                "_source": d
            }
            actions.append(action)
        
        # Bulk index documents
        if actions:
            from elasticsearch.helpers import bulk
            bulk(self.es_store.es, actions)
            
            # Force refresh to make documents immediately searchable
            self.es_store.es.indices.refresh(index=self.es_store.index_name)
    
    def remove_document(self, id_to_delete):
        """Remove a document by ID"""
        return self.es_store.remove_document(id_to_delete)
    
    def remove_source(self, source: str):
        """Remove documents by source"""
        return self.es_store.remove_source(source)
    
    def update_documents(self, doc_dicts: dict, **kwargs):
        """Update documents"""
        return self.es_store.update_documents(doc_dicts, **kwargs)
    
    def get_all_docs(self):
        """Get all documents"""
        return self.es_store.get_all_docs()
    
    def get_doc(self, id):
        """Get document by ID"""
        return self.es_store.get_doc(id)
    
    def get_size(self):
        """Get total number of documents"""
        return self.es_store.get_size()
    
    def erase(self, confirm=True):
        """Erase all documents"""
        return self.es_store.erase(confirm=confirm)
    
    def search(self, query: str, limit: int = 4, **kwargs):
        """Perform semantic search (alias for semantic_search)"""
        return self.semantic_search(query, limit=limit, return_dict=True, **kwargs)
    
    def semantic_search(self, query: str, limit: int = 4, 
                       filters: Optional[Dict[str, str]] = None,
                       return_dict: bool = False, **kwargs):
        """Perform semantic search using dense vectors"""
        if not hasattr(self, 'embeddings'):
            raise ValueError("Embeddings not initialized. Cannot perform semantic search.")
        
        # Generate query embedding
        query_embedding = self.embeddings.embed_query(query)
        
        # Use script_score to get actual cosine similarity scores
        script_query = {
            "size": limit,
            "query": {
                "script_score": {
                    "query": {"match_all": {}},  # Start with all documents
                    "script": {
                        "source": "cosineSimilarity(params.query_vector, '" + self.dense_vector_field + "') + 1.0",
                        "params": {
                            "query_vector": query_embedding
                        }
                    }
                }
            }
        }
        
        # Add filters if provided
        if filters:
            filter_clauses = []
            for k, v in filters.items():
                filter_type = "terms" if isinstance(v, list) else "term"
                filter_clauses.append({filter_type: {k: v}})
            script_query["query"]["script_score"]["query"] = {
                "bool": {
                    "filter": filter_clauses
                }
            }
        
        # Execute search
        response = self.es_store.es.search(index=self.es_store.index_name, body=script_query)
        
        # Process results
        hits = []
        for hit in response['hits']['hits']:
            doc_dict = hit['_source'].copy()
            # Convert script_score back to actual cosine similarity
            # script_score returns (cosineSimilarity + 1.0), so subtract 1.0
            elasticsearch_score = hit['_score']
            cosine_similarity = elasticsearch_score - 1.0
            doc_dict['score'] = cosine_similarity
            hits.append(doc_dict)
        
        total_hits = response['hits']['total']['value']
        
        if return_dict:
            return {'hits': hits, 'total_hits': total_hits}
        else:
            from ..helpers import doc_from_dict
            return [doc_from_dict(hit) for hit in hits]
    
    # get_db() method removed - use store methods instead

