"""full-text search engine as a backend vectorstore"""

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

# %% auto 0
__all__ = ['N_CANDIDATES', 'DEFAULT_SCHEMA', 'KEYWORD_ANALYZER', 'SparseStore', 'ReadOnlySparseStore', 'SharePointStore',
           'default_schema', 'WhooshStore', 'ElasticsearchSparseStore']

# %% ../../../nbs/01_ingest.stores.sparse.ipynb 3
from .base import VectorStore
from ..helpers import doc_from_dict
N_CANDIDATES = 12

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

    @classmethod
    def create(cls, persist_location=None, kind=None, **kwargs) -> 'SparseStore':
        """
        Factory method to construct a `SparseStore` instance.       
        Extra kwargs passed to object instantiation.
        
        Args:
            persist_location: where the index is stored (for whoosh) or Elasticsearch URL (for elasticsearch)
            kind: one of {whoosh, elasticsearch}
            
        Elasticsearch-specific kwargs:
            basic_auth: tuple of (username, password) for basic authentication
            verify_certs: whether to verify SSL certificates (default: True)
            ca_certs: path to CA certificate file
            timeout: connection timeout in seconds (default: 30, becomes request_timeout for v9+ compatibility)
            max_retries: maximum number of retries (default: 3)
            retry_on_timeout: whether to retry on timeout (default: True)
            maxsize: maximum number of connections in the pool (default: 25, removed for Elasticsearch v9+ compatibility)

        Returns:
            SparseStore instance
        """
        
        kind = 'whoosh' if not kind else kind
        
        if kind == 'whoosh':
            return WhooshStore(persist_location=persist_location, **kwargs)
        elif kind == 'elasticsearch':
            if not ELASTICSEARCH_INSTALLED:
                raise ImportError("Please install the elasticsearch package version for your "
                                  "Elasticsearch instance: e.g., pip install elasticsearch==9")
            return ElasticsearchSparseStore(persist_location=persist_location, **kwargs)
        else:
            raise ValueError(f"Unknown SparseStore type: {kind}")



    def _is_boolean_or_phrase_query(self, query: str) -> bool:
        """
        Returns True if the query looks like a Boolean or phrase query.
        """
        # Check for quoted phrases
        if re.search(r'"[^"]+"', query):
            return True

        # Check for Boolean operators (case-insensitive)
        if re.search(r'\b(AND|OR|NOT)\b', query, re.IGNORECASE):
            return True

        return False

    def augment_query(self, query: str) -> str:
        """
        Augments a natural language query with extracted noun phrases,
        unless the query already looks like a Boolean or phrase query.
        """
        from onprem.utils import extract_noun_phrases

        if self._is_boolean_or_phrase_query(query):
            return query  # Don't modify Boolean/phrase queries

        noun_phrases = extract_noun_phrases(query)
        if not noun_phrases:
            return query

        quoted_nps = [f'"{np}"^2.0' for np in noun_phrases]
        or_clause = " OR ".join(quoted_nps)

        return f"({query}) OR ({or_clause})"

    def _preprocess_query(self, query):
        """
        Removes question marks at the end of queries.
        This essentially disables using the question mark
        wildcard at end of search term so legitimate
        questions are not treated differntly depending
        on existence of question mark.
        """
        # Replace question marks at the end of the query
        if query.endswith('?'):
            query = query[:-1]

        # Handle quoted phrases with question marks at the end
        import re
        # Match question marks at the end of words or at the end of quoted phrases
        query = re.sub(r'(\w)\?([\s\"]|$)', r'\1\2', query)
        return query

    
    def get_content_field(self):
        """
        Get content field name (default to 'page_content')
        """
        return getattr(self, 'content_field', 'page_content')

        
    def semantic_search(self, *args, **kwargs):
        """
        Any subclass of SparseStore can inherit this method for on-the-fly semantic searches.
        Retrieves results based on semantic similarity to supplied `query`.
        All arguments are fowrwarded on to the store's query method.
        The query method is expected to include "query" as first positional argument and a "limit" keyword argument.
        Results of query method are expected to be in the form: {'hits': list_of_dicts, 'total_hits' : int}.
        Result of this method is a list of  LangChain Document objects sorted by semantic similarity.
        
        If subclass supports dynamic chunking (has chunk_for_semantic_search=True), 
        it will chunk large documents and find the best matching chunks per document.
        
        Args:
            return_chunks (bool): If True (default), return individual chunks as Document objects for RAG.
                                 If False, return original documents with full content and all chunk scores.
            load_web_documents (bool): If True, attempt to load content from web URLs when content field is empty (default: False).
            verbose (bool): If True (default), show progress bar during semantic search processing.
        """
        args = list(args)
        query = args[0] # for semantic search
        args[0] = self.augment_query(args[0]) # for keyword search
        args = tuple(args)

        limit = kwargs.get('limit', 4)
        n_candidates = kwargs.pop('n_candidates', getattr(self, 'n_candidates', limit * 10)) # check instance attribute first
        n_candidates = n_candidates if n_candidates and n_candidates>0 else limit*10 # if n_candidates is 0 or None
        return_chunks = kwargs.pop('return_chunks', True)  # Default True for RAG optimization
        load_web_documents = kwargs.pop('load_web_documents', getattr(self, 'load_web_documents', False))  # Check instance attribute first
        verbose = kwargs.pop('verbose', True)  # Progress bar enabled by default
        
        # Create a copy of kwargs to avoid modifying the original
        query_kwargs = kwargs.copy()
        query_kwargs['limit'] = n_candidates
        
        results = self.search(*args, **query_kwargs)['hits']
        
        # If no results but we have filters/where_document, try a broader search
        if not results and (kwargs.get('filters') or kwargs.get('where_document')):
            # Try with a wildcard query to get all documents, then filter
            wildcard_args = ('*',) + args[1:]  # Replace query with wildcard
            results = self.search(*wildcard_args, **query_kwargs)['hits']
        
        if not results: return []
        
        # Load web documents if enabled and content is empty
        if load_web_documents:
            results = self._load_web_documents_if_needed(results, verbose=verbose)
        
        # Check if subclass supports dynamic chunking
        if hasattr(self, 'chunk_for_semantic_search') and self.chunk_for_semantic_search:
            return self._semantic_search_with_chunking(query, results, limit, return_chunks, verbose=verbose)
        else:
            return self._semantic_search_original(query, results, limit, return_chunks, verbose=verbose)
    
    def _semantic_search_original(self, query, results, limit, return_chunks=True, verbose=True):
        """Original semantic search implementation without chunking.
        
        Args:
            return_chunks (bool): Included for API consistency. Since this method doesn't chunk,
                                 this parameter doesn't affect behavior.
        """
        texts = [r[self.get_content_field()] or '' for r in results]

        cos_scores = self.compute_similarity(query, texts)

        # Assign scores back to results
        for i, score in enumerate(cos_scores):
            results[i]['score'] = score

        # Sort results by similarity in descending order
        sorted_results = sorted(results, key=lambda x: x['score'], reverse=True)[:limit]
        return [doc_from_dict(r) for r in sorted_results]
    
    def _semantic_search_with_chunking(self, query, results, limit, return_chunks=True, verbose=True):
        """
        Enhanced semantic search that dynamically chunks large documents.
        Subclasses can override this method to customize chunking behavior.
        
        Args:
            return_chunks (bool): If True, return individual chunks as Document objects for RAG.
                                 If False, return original documents with full content and all chunk scores.
        """
        from ..base import chunk_documents
        
        content_field = self.get_content_field()
        chunk_size = getattr(self, 'chunk_size', 500)
        chunk_overlap = getattr(self, 'chunk_overlap', 50)
        
        # Process each document and its chunks
        doc_chunk_data = []
        all_chunks = []
        
        for doc_idx, doc in enumerate(results):
            # Extract text and use chunk_documents to split it
            doc_text = doc[content_field] or ''
            chunk_docs = chunk_documents(
                [doc_text],  # Pass text string directly
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap
            )
            
            # Store mapping of chunks to original document
            for chunk_idx, chunk_doc in enumerate(chunk_docs):
                doc_chunk_data.append({
                    'doc_idx': doc_idx,
                    'chunk_idx': chunk_idx,
                    'chunk_text': chunk_doc.page_content,
                    'original_doc': doc
                })
                all_chunks.append(chunk_doc.page_content)
        
        if not all_chunks:
            return []
        
        # Compute similarity for all chunks
        cos_scores = self.compute_similarity(query, all_chunks)
        
        # Store chunk scores and find best chunk for each document
        doc_best_scores = {}
        doc_best_chunks = {}
        chunk_score_data = []  # Store all chunks with their scores
        
        for i, chunk_data in enumerate(doc_chunk_data):
            doc_idx = chunk_data['doc_idx']
            score = cos_scores[i]
            
            # Store chunk with score for potential individual return
            chunk_score_data.append({
                'doc_idx': doc_idx,
                'chunk_idx': chunk_data['chunk_idx'],
                'chunk_text': chunk_data['chunk_text'],
                'score': score,
                'original_doc': chunk_data['original_doc']
            })
            
            # Keep track of best chunk per document
            if doc_idx not in doc_best_scores or score > doc_best_scores[doc_idx]:
                doc_best_scores[doc_idx] = score
                doc_best_chunks[doc_idx] = {
                    'chunk_text': chunk_data['chunk_text'],
                    'chunk_idx': chunk_data['chunk_idx'],
                    'score': score
                }
        
        if return_chunks:
            # Return individual chunks as Document objects, sorted by score
            chunk_score_data.sort(key=lambda x: x['score'], reverse=True)
            chunk_score_data = chunk_score_data[:limit]
            
            documents = []
            for chunk_data in chunk_score_data:
                # Create metadata from original document and add chunk-specific info
                metadata = chunk_data['original_doc'].copy()
                metadata['score'] = chunk_data['score']
                metadata['chunk_idx'] = chunk_data['chunk_idx']
                metadata['is_chunk'] = True
                
                doc = Document(
                    page_content=chunk_data['chunk_text'],
                    metadata=metadata
                )
                documents.append(doc)
            
            return documents
        else:
            # Return original documents with scores from best chunks (backward compatibility)
            final_results = []
            for doc_idx, best_chunk in doc_best_chunks.items():
                doc = results[doc_idx].copy()
                doc['score'] = best_chunk['score']
                
                # Get all chunks for this document with their scores
                doc_chunks_with_scores = []
                for chunk_data in chunk_score_data:
                    if chunk_data['doc_idx'] == doc_idx:
                        doc_chunks_with_scores.append({
                            'text': chunk_data['chunk_text'],
                            'idx': chunk_data['chunk_idx'],
                            'score': chunk_data['score']
                        })
                
                # Sort chunks by their original order (by chunk_idx)
                doc_chunks_with_scores.sort(key=lambda x: x['idx'])
                
                # Extract parallel arrays
                doc_chunks = [item['text'] for item in doc_chunks_with_scores]
                doc_chunk_scores = [item['score'] for item in doc_chunks_with_scores]
                
                # Add metadata about all chunks and the best matching chunk
                doc['best_chunk_text'] = best_chunk['chunk_text']
                doc['best_chunk_idx'] = best_chunk['chunk_idx']
                doc['total_chunks'] = len(doc_chunks)
                doc['chunk_scores'] = doc_chunk_scores  # Parallel list of scores
                
                final_results.append((doc, doc_chunks))
            
            # Sort by best chunk score and limit
            final_results.sort(key=lambda x: x[0]['score'], reverse=True)
            final_results = final_results[:limit]
            
            # Create Document objects with chunks as content
            documents = []
            for doc_data, chunks in final_results:
                # Store original chunks list in metadata
                doc_data['chunks'] = chunks
                doc = Document(
                    page_content='\n\n'.join(chunks),  # Join chunks for LangChain compatibility
                    metadata=doc_data
                )
                documents.append(doc)
            
            return documents
    
    def _load_web_documents_if_needed(self, results, verbose=True):
        """
        Load content from web URLs when content field is empty.
        Only processes URLs that start with 'http://' or 'https://'.
        """
        from ..base import load_single_document
        import tempfile
        import os
        import requests
        from urllib.parse import urlparse
        
        from tqdm import tqdm
        
        content_field = self.get_content_field()
        updated_results = []
        
        # Count documents that need web loading for progress bar
        web_docs_needed = sum(1 for result in results 
                             if not result.get(content_field, '').strip() 
                             and result.get('source', '').strip() 
                             and self._is_web_url(result.get('source', '')))
        
        if web_docs_needed > 0 and verbose:
            progress_bar = tqdm(results, desc=f"Loading {web_docs_needed} web documents", total=len(results))
        else:
            progress_bar = results
        
        for result in progress_bar:
            # Check if content is empty or missing
            content = result.get(content_field, '').strip()
            source = result.get('source', '').strip()
            
            if not content and source and self._is_web_url(source):
                try:
                    # Download the document from the web URL
                    loaded_docs = self._load_web_document(source)
                    if loaded_docs:
                        # Concatenate content from all pages/documents
                        result[content_field] = '\n\n'.join([doc.page_content for doc in loaded_docs])
                        # Optionally merge metadata
                        if hasattr(loaded_docs[0], 'metadata'):
                            for key, value in loaded_docs[0].metadata.items():
                                if key not in result:
                                    result[key] = value
                except Exception as e:
                    # Log warning but continue with empty content
                    import warnings
                    warnings.warn(f"Failed to load web document from {source}: {str(e)}")
            
            updated_results.append(result)
        
        return updated_results
    
    def _is_web_url(self, url):
        """
        Check if the given URL is a web URL (starts with http:// or https://).
        """
        return url.startswith(('http://', 'https://'))
    
    def _load_web_document(self, url):
        """
        Download and extract text from a web document using load_single_document.
        """
        from ..base import load_web_document
        
        # Get credentials from SharePointStore instance if available
        username = getattr(self, 'username', None)
        password = getattr(self, 'password', None)
        
        return load_web_document(url, username=username, password=password)

# %% ../../../nbs/01_ingest.stores.sparse.ipynb 6
class ReadOnlySparseStore(SparseStore):
    """
    A sparse vector store based on a read-only full-text search engine
    """

    def exists(self):
        return True


    def add_documents(self, documents):
        """
        Adds documents.
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')
        
    def remove_document(self, id_to_delete):
        """
        Remove a single document.
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')

    
    def remove_source(self, source):
        """
        Remove documents by source
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')

    def update_documents(self, *args, **kwargs):
        """
        Update a set of documents.
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')

    def get_all_docs(self):
        """
        Returns a list of files previously added to vector store.
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')        

    def get_size(self):
        """
        Get total number of records added to vector store
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')

    def erase(self):
        """
        Removes all documents in vector store
        """
        raise NotImplementedError('Not supported for ReadOnlySparseStore instances.')

# %% ../../../nbs/01_ingest.stores.sparse.ipynb 7
from typing import Optional, Dict
import requests
try:
    from requests_ntlm import HttpNtlmAuth
    REQUESTS_NTLM_INSTALLED = True
except ImportError:
    REQUESTS_NTLM_INSTALLED = False
    
    
class SharePointStore(ReadOnlySparseStore):
    """
    A sparse vector store based on Microsoft Sharepoint using a "vectors-on-demand" approach.
    """
    def __init__(self,
                persist_location:str,  # SharePoint URL
                username:str, # user name (e.g., CORP\my_username)
                password:str, # your SharePoint password
                chunk_for_semantic_search: bool = True, # whether to do dynamic chunking
                chunk_size: int = 500, # chunk size in characters (1 word ~ 4 characters)
                chunk_overlap: int = 50, # chunk overlap between chunks
                n_candidates:Optional[int]=None, # number of candidate documents to consider for answer. Defaults to limit*10.
                load_web_documents: bool = True, # whether to load content from web URLs when content field is empty
                **kwargs, 
        ):
        """
        Initializes full-text search engine based on Microsoft Sharepoint
        """

        if not REQUESTS_NTLM_INSTALLED:
            raise ImportError('Please install requests_ntlm: pip install requests_ntlm')

        self.persist_location = persist_location # SharePoint URL
        self.username = username
        self.password = password
        self.chunk_for_semantic_search = chunk_for_semantic_search
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        if n_candidates: self.n_candidates = n_candidates
        self.load_web_documents = load_web_documents
        self.init_embedding_model(**kwargs)
        

    def search(self, 
               query:str,  # query string
               is_document:int=1,   # search documents by default
               filters:Optional[Dict[str,str]]=None,  # filters
               where_document:Optional[str]=None,     # boolean clause appended to query
               limit=10,  # rowlimit in SharePoint
               page=0,    # startrow in SharePoint
               trimduplicates=False, # Prevents SharePoint from collapsing results it thinks are duplicates. This often changes total count.
               enablequeryrules=False, # Turns off "smart" query rewriting that SharePoint applies for predictability
               **kwargs):
        """
        Perform SharePoint keyword search for documents
        
        Args:
            query: search query string
            is_document: filter for document type (default: 1)
            filters: dictionary of field filters (e.g., {'FileExtension': 'pdf'})
            where_document: additional query to filter documents
        """
        # Start with base query
        search_query = f"({query}) AND (IsDocument={str(is_document)})"
        
        # Add filters
        if filters:
            for field, value in filters.items():
                search_query += f" AND ({field}:{value})"
        
        # Add where_document filter
        if where_document:
            search_query += f" AND ({where_document})"
        
        from urllib.parse import quote
        encoded_query = quote(search_query)
        search_url = (
            f"{self.persist_location}/_api/search/query"
            f"?querytext='{encoded_query}'"
            f"&selectproperties='Title,Path,FileExtension,Contents'"
            f"&rowlimit={limit}&startrow={page}"
            f"&trimduplicates={'true' if trimduplicates else 'false'}"
            f"&enablequeryrules={'true' if enablequeryrules else 'false'}"
        )

        # Make request with NTLM (Windows auth)
        response = requests.get(
            search_url,
            auth=HttpNtlmAuth(self.username, self.password),
            headers={"Accept": "application/json;odata=verbose"}
        )
        

        # Parse and print results
        total_hits = -1
        if response.status_code == 200:
            results = response.json()
            rows = results['d']['query']['PrimaryQueryResult']['RelevantResults']['Table']['Rows']['results']
            total_hits = results['d']['query']['PrimaryQueryResult']['RelevantResults']['TotalRows']

            results = []
            for row in rows:
                props = {item['Key']: item['Value'] for item in row['Cells']['results']}
                title = props.get('Title', '')
                path = props.get('Path', '')
                contents = props.get('Contents', '') or ''
                id = props.get('DocId', '')
                results.append({'id':id,
                                'source': path,
                                'title': title,
                                'page_content': contents})
            return {'hits':results, 'total_hits':total_hits}
        else:
            response.raise_for_status()

    def get_doc(self, id: str):
        """
        Get a specific document by ID from SharePoint
        """
        query = f"DocId:{id}"
        results = self.search(query)
        return results['hits'][0] if results['hits'] else None

# %% ../../../nbs/01_ingest.stores.sparse.ipynb 9
import json
import os
import warnings
from typing import Dict, List, Optional, Sequence
import math
import numpy as np

from whoosh import index
from whoosh.analysis import StemmingAnalyzer
from whoosh.fields import *
from whoosh.filedb.filestore import RamStorage
from whoosh.qparser import MultifieldParser, OrGroup
from whoosh.query import Term, And, Variations
from whoosh.analysis import RegexTokenizer, LowercaseFilter
from langchain_core.documents import Document
import uuid
from tqdm import tqdm
from sentence_transformers import SentenceTransformer

try:
    from elasticsearch import Elasticsearch
    ELASTICSEARCH_INSTALLED = True
except ImportError:
    ELASTICSEARCH_INSTALLED = False

from ..helpers import doc_from_dict

# ------------------------------------------------------------------------------
# IMPORTANT: The schema below contains only essential fields required for core
#            functionality. All other fields (filepath, filename, tags, etc.) are
#            automatically added as dynamic fields when documents are indexed.
#
#            The page_content field is the only truly required field in supplied
#            Document objects. All other fields, including dynamic fields, are optional. 
# ------------------------------------------------------------------------------

DEFAULT_SCHEMA = Schema(
    page_content=TEXT(stored=True, analyzer=StemmingAnalyzer()), # REQUIRED with stemming enabled
    id=ID(stored=True, unique=True),
    source=KEYWORD(stored=True, commas=True), 
    source_search=TEXT(stored=True, analyzer=StemmingAnalyzer()),
    )
DEFAULT_SCHEMA.add("*_t", TEXT(stored=True, analyzer=StemmingAnalyzer()), glob=True)
DEFAULT_SCHEMA.add("*_k", KEYWORD(stored=True, commas=True), glob=True)
DEFAULT_SCHEMA.add("*_b", BOOLEAN(stored=True), glob=True)
DEFAULT_SCHEMA.add("*_n", NUMERIC(stored=True), glob=True)
DEFAULT_SCHEMA.add("*_d", DATETIME(stored=True), glob=True)


def default_schema():
    schema = DEFAULT_SCHEMA
    #if "raw" not in schema.stored_names():
        #schema.add("raw", TEXT(stored=True))
    return schema

KEYWORD_ANALYZER = RegexTokenizer() | LowercaseFilter()

class WhooshStore(SparseStore):
    """
    A sparse vector store based on the Whoosh full-text search engine.
    """
    def __init__(self,
                persist_location: Optional[str]=None, 
                index_name:str = 'myindex',
                # Dynamic chunking for semantic search
                chunk_for_semantic_search: bool = False,
                chunk_size: int = 500,
                chunk_overlap: int = 50,
                **kwargs,
        ):
        """
        Initializes full-text search engine.

        **Args:**

        - *persist_location*: path to folder where search index is stored
        - *index_name*: name of index
        - *chunk_for_semantic_search*: if True, dynamically chunk large documents for semantic search (default: False)
        - *chunk_size*: size of chunks when chunk_for_semantic_search=True (default: 500)
        - *chunk_overlap*: overlap between chunks when chunk_for_semantic_search=True (default: 50)
        - *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}`).
        """

        self.persist_location = persist_location # alias for consistency with DenseStore
        self.index_name = index_name
        
        # Store dynamic chunking settings
        self.chunk_for_semantic_search = chunk_for_semantic_search
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.content_field = 'page_content'  # WhooshStore uses standard field name
        if self.persist_location and not self.index_name:
            raise ValueError('index_name is required if persist_location is supplied')
        if self.persist_location:
            if not index.exists_in(self.persist_location, indexname=self.index_name):
                self.ix = __class__.initialize_index(self.persist_location, self.index_name)
            else:
                self.ix = index.open_dir(self.persist_location, indexname=self.index_name)
        else:
            warnings.warn(
                "No persist_location was supplied, so an in-memory only index "
                "was created using DEFAULT_SCHEMA"
            )
            self.ix = RamStorage().create_index(default_schema())
        self.init_embedding_model(**kwargs) # stored as self.embeddings

    @classmethod
    def index_exists_in(cls, index_path: str, index_name: Optional[str] = None):
        """
        Returns True if index exists with name, *indexname*, and path, *index_path*.
        """
        return index.exists_in(index_path, indexname=index_name)

    @classmethod
    def initialize_index(
        cls, index_path: str, index_name: str, schema: Optional[Schema] = None
    ):
        """
        Initialize index

        **Args**

        - *index_path*: path to folder storing search index
        - *index_name*: name of index
        - *schema*: optional whoosh.fields.Schema object.
                    If None, DEFAULT_SCHEMA is used
        """
        schema = default_schema() if not schema else schema

        if index.exists_in(index_path, indexname=index_name):
            raise ValueError(
                f"There is already an existing index named {index_name}  with path {index_path} \n"
                + f"Delete {index_path} manually and try again."
            )
        if not os.path.exists(index_path):
            os.makedirs(index_path)
        ix = index.create_in(index_path, indexname=index_name, schema=schema)
        return ix


    def doc2dict(self, doc:Document):
        """
        Convert LangChain Document to expected format
        """
        d = {}
        for k,v in doc.metadata.items():
            d[k] = v
        d['id'] = uuid.uuid4().hex if not doc.metadata.get('id', '') else doc.metadata['id']
        d['page_content' ] = self.normalize_text(doc.page_content)
        #d['raw'] = json.dumps(d)
        if 'source' in d:
            d['source_search'] = d['source']
        if 'filepath' in d:
            d['filepath_search'] = d['filepath']
        return d

    def delete_by_prefix(self, prefix:str, field:str, optimize:bool=False):
        """
        Deletes all documents from a Whoosh index where the `source_field` starts with the given prefix.

        **Args:**

        - *prefix*: The prefix to match in the `field`.
        - *field*: The name of the field to match against.
        - *optimize*: If True, optimize when committing.

        **Returns:**
        - Number of records deleted
		"""

        from whoosh.query import Prefix
        with self.ix.searcher() as searcher:
            results = searcher.search(Prefix(field, prefix), limit=None)

            if results:
                writer = self.ix.writer()
                for hit in results:
                    writer.delete_document(hit.docnum)
                writer.commit(optimize=optimize)
                return len(results)
            else:
                return 0

    def close_index(self):
        """
        Run this if there are lock errors when trying to write to index
        """
        if hasattr(self, 'ix') and self.ix:
            self.ix.close()
    
    #------------------------------
    # overrides of abstract methods
    # -----------------------------


    def exists(self):
        """
        Returns True if documents have been added to search index
        """
        return self.get_size() > 0


    def add_documents(self,
                      docs: Sequence[Document], # list of LangChain Documents
                      limitmb:int=1024, # maximum memory in  megabytes to use
                      optimize:bool=False, # whether or not to also opimize index
                      verbose:bool=True, # Set to False to disable progress bar
                      **kwargs,
        ):
        """
        Indexes documents. Extra kwargs supplied to `TextStore.ix.writer`.
        """
        writer = self.ix.writer(limitmb=limitmb, **kwargs)
        
        # Add any new fields first
        stored_names = self.ix.schema.stored_names()
        new_fields = set()
        
        for doc in docs:
            for k, v in doc.metadata.items():
                if k not in stored_names and k not in new_fields:
                    # Add field to writer based on type
                    if isinstance(v, bool):
                        writer.add_field(k, BOOLEAN(stored=True))
                    elif isinstance(v, list):
                        # Lists are stored as comma-separated keywords
                        #writer.add_field(k, KEYWORD(stored=True, commas=True))
                        writer.add_field(k, KEYWORD(stored=True, commas=True, analyzer=KEYWORD_ANALYZER))
                    elif isinstance(v, str):
                        if k.endswith('_date'):
                            writer.add_field(k, DATETIME(stored=True))
                        elif len(v) > 100:  # Use TEXT for strings longer than 100 characters
                            writer.add_field(k, TEXT(stored=True, analyzer=StemmingAnalyzer()))
                        else:
                            writer.add_field(k, KEYWORD(stored=True, commas=True, analyzer=KEYWORD_ANALYZER))
                    elif isinstance(v, (int, float)):
                        writer.add_field(k, NUMERIC(stored=True))
                    new_fields.add(k)
        
        for doc in tqdm(docs, total=len(docs), disable=not verbose):
            d = self.doc2dict(doc)
            writer.update_document(**d)
        writer.commit(optimize=optimize)

       
    def remove_document(self, value:str, field:str='id', optimize:bool=False):
        """
        Remove document with corresponding value and field.
        Default field is the id field.
        If optimize is True, index will be optimized.
        """
        writer = self.ix.writer()
        writer.delete_by_term(field, value)
        writer.commit(optimize=optimize)
        return


    def remove_source(self, source:str, optimize:bool=False):
        """
        remove all documents associated with `source`.
        The `source` argument can either be the full path to
        document or a parent folder.  In the latter case,
        ALL documents in parent folder will be removed.
        If optimize is True, index will be optimized.
        """
        return self.delete_by_prefix(source, field='source', optimize=optimize)
        

    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)
        """
        docs = [doc_from_dict(d) for d in doc_dicts]
        self.add_documents(docs)


    def get_all_docs(self):
        """
        Returns a generator to iterate through all indexed documents
        """
        return self.ix.searcher().documents()
       

    def get_doc(self, id:str):
        """
        Get an indexed record by ID
        """
        r = self.search(f'id:{id}', return_dict=True)
        return r['hits'][0] if len(r['hits']) > 0 else None


    def get_size(self, include_deleted:bool=False) -> int:
        """
        Gets size of index

        If include_deleted is True, will include deletd detects (prior to optimization).
        """
        return self.ix.doc_count_all() if include_deleted else self.ix.doc_count()

        
    def erase(self, confirm=True):
        """
        Clears index
        """
        shall = True
        if confirm:
            msg = (
                f"You are about to remove all documents from the search index."
                + f"(Original documents on file system will remain.) Are you sure?"
            )
            shall = input("%s (Y/n) " % msg) == "Y"
        if shall and index.exists_in(
            self.persist_location, indexname=self.index_name
        ):
            ix = index.create_in(
                self.persist_location,
                indexname=self.index_name,
                schema=default_schema(),
            )
            return True
        return False


    def search(
            self,
            query: str,
            fields: Sequence = ["page_content"],
            highlight: bool = True,
            limit:int=10,
            page:int=1,
            return_dict:bool=True,
            filters:Optional[Dict[str, str]] = None,
            where_document:Optional[str]=None,
            return_generator=False,
            **kwargs

    ):
        """
        Queries the index

        **Args**

        - *query*: the query string
        - *fields*: a list of fields to search
        - *highlight*: If True, highlight hits
        - *limit*: results per page
        - *page*: page of hits to return
        - *return_dict*: If True, return list of dictionaries instead of LangChain Document objects
        - *filters*: filter results by field values (e.g., {'extension':'pdf'})
        - *where_document*: optional query to further filter results
        - *return_generator*: If True, returns a generator, not a list
        """

        q = query
        q = self._preprocess_query(q)

        if where_document:
            q = f'({q}) AND ({where_document})'

        # process filters
        combined_filter=None
        if filters:
            terms = []
            for k in filters:
                terms.append(Term(k, filters[k]))
            combined_filter = And(terms)
        
        def get_search_results(searcher):
            """Helper to get search results from searcher"""
            if page == 1:
                return searcher.search(
                    MultifieldParser(fields, schema=self.ix.schema, termclass=Variations, group=OrGroup.factory(0.9)).parse(q), limit=limit, filter=combined_filter)
            else:
                return searcher.search_page(
                    MultifieldParser(fields, schema=self.ix.schema, termclass=Variations, group=OrGroup.factory(0.9)).parse(q), page, limit, filter=combined_filter)
        
        def process_result(r):
            """Helper to process individual search result"""
            d = dict(r)
            if highlight:
                for f in fields:
                    if r[f] and isinstance(r[f], str):
                        d['hl_'+f] = r.highlights(f) or r[f]
            return d if return_dict else doc_from_dict(d)
        
        if return_generator:
            searcher = self.ix.searcher()
            try:
                results = get_search_results(searcher)
                total_hits = results.scored_length()
                if page > math.ceil(total_hits/limit):
                   results = []
                
                def result_generator():
                    try:
                        for r in results:
                            yield process_result(r)
                    finally:
                        searcher.close()
                
                return {'hits': result_generator(),
                        'total_hits': total_hits}
            except:
                searcher.close()
                raise
        else:
            with self.ix.searcher() as searcher:
                results = get_search_results(searcher)
                total_hits = results.scored_length()
                if page > math.ceil(total_hits/limit):
                   results = []
                
                return {'hits': [process_result(r) for r in results],
                        'total_hits': total_hits}


class ElasticsearchSparseStore(SparseStore):
    """
    A sparse vector store based on Elasticsearch.
    """
    def __init__(self,
                persist_location: Optional[str]=None,
                index_name:str = 'myindex',
                basic_auth: Optional[tuple] = None,
                verify_certs: bool = True,
                ca_certs: Optional[str] = None,
                timeout: int = 30,
                max_retries: int = 3,
                retry_on_timeout: bool = True,
                maxsize: int = 25,
                # Field mapping parameters for existing indices
                content_field: str = 'page_content',
                source_field: Optional[str] = 'source',
                id_field: str = 'id',
                content_analyzer: str = 'standard',
                # Dynamic chunking for semantic search
                chunk_for_semantic_search: bool = False,
                chunk_size: int = 500,
                chunk_overlap: int = 50,
                n_candidates:Optional[int]=None,
                **kwargs,
        ):
        """
        Initializes Elasticsearch-based full-text search engine.

        **Args:**

        - *persist_location*: Elasticsearch URL (e.g., 'http://localhost:9200'). If None, defaults to 'http://localhost:9200'
        - *index_name*: name of Elasticsearch index
        - *basic_auth*: tuple of (username, password) for basic authentication
        - *verify_certs*: whether to verify SSL certificates
        - *ca_certs*: path to CA certificate file
        - *timeout*: connection timeout in seconds (becomes request_timeout internally for Elasticsearch v9+ compatibility)
        - *max_retries*: maximum number of retries
        - *retry_on_timeout*: whether to retry on timeout
        - *maxsize*: maximum number of connections in the pool (removed for Elasticsearch v9+ compatibility)
        - *content_field*: field name for document content (default: 'page_content')
        - *source_field*: field name for document source (default: 'source', set to None to disable)
        - *id_field*: field name for document ID (default: 'id')
        - *content_analyzer*: analyzer for content field (default: 'standard')
        - *chunk_for_semantic_search*: if True, dynamically chunk large documents for semantic search (default: False)
        - *chunk_size*: size of chunks when chunk_for_semantic_search=True (default: 500)
        - *chunk_overlap*: overlap between chunks when chunk_for_semantic_search=True (default: 50)
        - *n_candidates*:  number of candidate documents to consider for answer. Defaults to n*10
        - *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}`).
        """
        if not ELASTICSEARCH_INSTALLED:
            raise ImportError('Please install a version of elasticsearch compatible with your running Elasticsearch instance. For latest: pip install elasticsearch')

        # Use persist_location as Elasticsearch URL
        self.elasticsearch_url = persist_location if persist_location else 'http://localhost:9200'
        self.persist_location = self.elasticsearch_url  # for interface compatibility
        self.index_name = index_name
        
        # Store field mappings for custom field names
        self.content_field = content_field
        self.source_field = source_field
        self.id_field = id_field
        self.content_analyzer = content_analyzer
        
        # Store dynamic chunking settings
        self.chunk_for_semantic_search = chunk_for_semantic_search
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        if n_candidates: self.n_candidates = n_candidates
        
        # Prepare Elasticsearch client parameters
        # Use request_timeout (works in both v8 and v9) instead of timeout (removed in v9)
        # Remove maxsize (removed in v9) for compatibility
        es_params = {
            'request_timeout': timeout,
            'max_retries': max_retries,
            'retry_on_timeout': retry_on_timeout,
        }
        
        # Add authentication if provided
        if basic_auth:
            es_params['basic_auth'] = basic_auth
            
        # Add SSL parameters
        if not verify_certs:
            es_params['verify_certs'] = False
        if ca_certs:
            es_params['ca_certs'] = ca_certs
            
        # Filter out embedding-related kwargs before passing to Elasticsearch
        embedding_kwargs = {}
        for k, v in kwargs.items():
            if k.startswith('embedding_'):
                embedding_kwargs[k] = v
            else:
                es_params[k] = v
        
        # Initialize Elasticsearch client
        self.es = Elasticsearch([self.elasticsearch_url], **es_params)
        
        # Handle index creation or validation
        if not self.es.indices.exists(index=self.index_name):
            self._create_index()
        else:
            # Validate existing index has required fields
            self._validate_existing_index()
        
        self.init_embedding_model(**embedding_kwargs)  # stored as self.embeddings

    def _create_index(self):
        """Create Elasticsearch index with appropriate mapping using custom field names and analyzer"""
        properties = {
            # Essential fields for core functionality using custom field names and analyzer
            self.content_field: {"type": "text", "analyzer": self.content_analyzer},
            self.id_field: {"type": "keyword"},
        }
        
        # Add source field only if specified
        if self.source_field:
            properties[self.source_field] = {"type": "keyword"}
            properties[f"{self.source_field}_search"] = {"type": "text", "analyzer": self.content_analyzer}
        
        mapping = {
            "mappings": {
                "properties": properties
            }
        }
        
        self.es.indices.create(index=self.index_name, body=mapping)

    def _validate_existing_index(self):
        """Validate that existing index has required fields with compatible types"""
        try:
            # Get index mapping
            mapping_response = self.es.indices.get_mapping(index=self.index_name)
            index_mapping = mapping_response[self.index_name]['mappings']
            properties = index_mapping.get('properties', {})
            
            # Check required fields exist
            missing_fields = []
            incompatible_fields = []
            
            # Validate content field
            if self.content_field not in properties:
                missing_fields.append(f"Content field '{self.content_field}' not found")
            else:
                field_props = properties[self.content_field]
                if field_props.get('type') != 'text':
                    incompatible_fields.append(f"Content field '{self.content_field}' is not text type (found: {field_props.get('type')})")
            
            # Validate ID field  
            if self.id_field not in properties:
                missing_fields.append(f"ID field '{self.id_field}' not found")
            else:
                field_props = properties[self.id_field]
                if field_props.get('type') not in ['keyword', '_id']:
                    incompatible_fields.append(f"ID field '{self.id_field}' is not keyword type (found: {field_props.get('type')})")
            
            # Validate source field (only if specified)
            if self.source_field:
                if self.source_field not in properties:
                    missing_fields.append(f"Source field '{self.source_field}' not found")
                else:
                    field_props = properties[self.source_field]
                    if field_props.get('type') not in ['keyword', 'text']:
                        incompatible_fields.append(f"Source field '{self.source_field}' is not keyword or text type (found: {field_props.get('type')})")
            
            # Report validation results
            if missing_fields:
                warnings.warn(f"Index validation warnings - Missing fields: {'; '.join(missing_fields)}")
            
            if incompatible_fields:
                warnings.warn(f"Index validation warnings - Incompatible field types: {'; '.join(incompatible_fields)}")
                
            # Log analyzer info for content field if it exists
            if self.content_field in properties:
                content_props = properties[self.content_field]
                existing_analyzer = content_props.get('analyzer', 'default')
                if existing_analyzer != self.content_analyzer:
                    warnings.warn(f"Content field '{self.content_field}' uses analyzer '{existing_analyzer}' but you specified '{self.content_analyzer}'. Search behavior may be affected.")
                    
        except Exception as e:
            warnings.warn(f"Could not validate index mapping: {e}")

    @classmethod
    def index_exists_in(cls, index_path: str, index_name: Optional[str] = None, elasticsearch_url: str = 'http://localhost:9200'):
        """
        Returns True if index exists in Elasticsearch cluster.
        
        **Args:**
        - *index_path*: Not used for Elasticsearch (kept for interface compatibility)
        - *index_name*: name of Elasticsearch index
        - *elasticsearch_url*: Elasticsearch URL (e.g., 'http://localhost:9200')
        """
        if not ELASTICSEARCH_INSTALLED:
            raise ImportError("Please install the elasticsearch package version for your "
                              "Elasticsearch instance: e.g., pip install elasticsearch==9)")
        
        es = Elasticsearch([elasticsearch_url])
        return es.indices.exists(index=index_name)

    @classmethod
    def initialize_index(cls, index_path: str, index_name: str, elasticsearch_url: str = 'http://localhost:9200'):
        """
        Initialize Elasticsearch index
        
        **Args:**
        - *index_path*: Not used for Elasticsearch (kept for interface compatibility) 
        - *index_name*: name of Elasticsearch index
        - *elasticsearch_url*: Elasticsearch URL (e.g., 'http://localhost:9200')
        """
        store = cls(persist_location=elasticsearch_url, index_name=index_name)
        return store.es


    def doc2dict(self, doc: Document):
        """
        Convert LangChain Document to expected format using custom field mappings
        """
        d = {}
        for k, v in doc.metadata.items():
            d[k] = v
        
        # Use custom field mappings
        d[self.id_field] = uuid.uuid4().hex if not doc.metadata.get('id', '') else doc.metadata['id']
        d[self.content_field] = self.normalize_text(doc.page_content)
        
        # Handle source field mapping (only if source field is configured)
        if self.source_field and 'source' in d:
            d[self.source_field] = d.pop('source')  # Replace 'source' with custom field name
            d[f"{self.source_field}_search"] = d[self.source_field]
        if 'filepath' in d:
            d['filepath_search'] = d['filepath']
        return d

    def delete_by_prefix(self, prefix: str, field: str):
        """
        Deletes all documents from Elasticsearch index where the field starts with the given prefix.
        """
        query = {
            "query": {
                "prefix": {
                    field: prefix
                }
            }
        }
        
        response = self.es.delete_by_query(index=self.index_name, body=query)
        return response.get('deleted', 0)


    # overrides of abstract methods
    def exists(self):
        """
        Returns True if documents have been added to search index
        """
        return self.get_size() > 0

    def add_documents(self,
                      docs: Sequence[Document],
                      limitmb: int = 1024,  # Not used in Elasticsearch
                      optimize: bool = False,  # Not used in Elasticsearch
                      verbose: bool = True,
                      **kwargs,
        ):
        """
        Indexes documents in Elasticsearch
        """
        actions = []
        for doc in tqdm(docs, total=len(docs), disable=not verbose):
            d = self.doc2dict(doc)
            action = {
                "_index": self.index_name,
                "_id": d['id'],
                "_source": d
            }
            actions.append(action)
        
        if actions:
            from elasticsearch.helpers import bulk
            bulk(self.es, actions)
            
            # Force refresh to make documents immediately searchable
            self.es.indices.refresh(index=self.index_name)
            
        if optimize:
            self.es.indices.forcemerge(index=self.index_name)

    def remove_document(self, value: str, field: str = 'id', optimize: bool = False):
        """
        Remove document with corresponding value and field.
        """
        query = {
            "query": {
                "term": {
                    field: value
                }
            }
        }
        
        self.es.delete_by_query(index=self.index_name, body=query)
        
        if optimize:
            self.es.indices.forcemerge(index=self.index_name)

    def remove_source(self, source: str, optimize: bool = False):
        """
        remove all documents associated with `source`.
        """
        if not self.source_field:
            raise ValueError("Cannot remove by source: no source field configured. Set source_field parameter when creating the store.")
        return self.delete_by_prefix(source, field=self.source_field)

    def update_documents(self,
                         doc_dicts: dict,
                         **kwargs):
        """
        Update a set of documents (doc in index with same ID will be over-written)
        """
        docs = [doc_from_dict(d) for d in doc_dicts]
        self.add_documents(docs)

    def get_all_docs(self):
        """
        Returns a generator to iterate through all indexed documents
        """
        query = {"query": {"match_all": {}}}
        response = self.es.search(index=self.index_name, body=query, scroll='1m', size=1000)
        
        while response['hits']['hits']:
            for hit in response['hits']['hits']:
                yield hit['_source']
            
            # Get next batch
            response = self.es.scroll(scroll_id=response['_scroll_id'], scroll='1m')

    def get_doc(self, id: str):
        """
        Get an indexed record by ID
        """
        try:
            response = self.es.get(index=self.index_name, id=id)
            return response['_source']
        except:
            return None

    def get_size(self, include_deleted: bool = False) -> int:
        """
        Gets size of index
        """
        response = self.es.count(index=self.index_name)
        return response['count']

    def erase(self, confirm=True):
        """
        Clears index
        """
        shall = True
        if confirm:
            msg = (
                f"You are about to remove all documents from the search index."
                + f"(Original documents on file system will remain.) Are you sure?"
            )
            shall = input("%s (Y/n) " % msg) == "Y"
        
        if shall:
            try:
                self.es.indices.delete(index=self.index_name)
                self._create_index()
                return True
            except:
                return False
        return False

    def search(self,
              query: str,
              fields: Sequence = None,
              highlight: bool = True,
              limit: int = 10,
              page: int = 1,
              return_dict: bool = True,
              filters: Optional[Dict[str, str]] = None,
              where_document: Optional[str] = None,
              return_generator=False,
              **kwargs
        ):
        """
        Queries the Elasticsearch index
        """
        # Use custom content field as default if no fields specified
        if fields is None:
            fields = [self.content_field]
        
        q = self._preprocess_query(query)
        
        es_query = {
            "query_string": {
                "query": q,
                "fields": fields,
                "default_operator": "OR"
            }
        }
        
        # Add additional document filter
        if where_document:
            es_query = {
                "bool": {
                    "must": [
                        es_query,
                        {"query_string": {"query": where_document}}
                    ]
                }
            }
        
        # Add filters
        if filters:
            filter_clauses = []
            for k, v in filters.items():
                filter_clauses.append({"term": {k: v}})
            
            if "bool" not in es_query:
                es_query = {"bool": {"must": [es_query]}}
            
            es_query["bool"]["filter"] = filter_clauses
        
        # Calculate from parameter for pagination
        from_param = (page - 1) * limit
        
        # Build search body
        search_body = {
            "query": es_query,
            "from": from_param,
            "size": limit
        }
        
        # Add highlighting if requested
        if highlight:
            search_body["highlight"] = {
                "fields": {field: {} for field in fields}
            }
        
        def process_result(hit):
            """Helper to process individual search result"""
            d = hit['_source'].copy()
            if highlight and 'highlight' in hit:
                for field in fields:
                    if field in hit['highlight']:
                        d[f'hl_{field}'] = ' '.join(hit['highlight'][field])
                    else:
                        d[f'hl_{field}'] = d.get(field, '')
            return d if return_dict else doc_from_dict(d)
        
        if return_generator:
            def result_generator():
                response = self.es.search(index=self.index_name, body=search_body)
                for hit in response['hits']['hits']:
                    yield process_result(hit)
            
            # Get total count
            count_response = self.es.count(index=self.index_name, body={"query": es_query})
            total_hits = count_response['count']
            
            return {'hits': result_generator(), 'total_hits': total_hits}
        else:
            response = self.es.search(index=self.index_name, body=search_body)
            hits = [process_result(hit) for hit in response['hits']['hits']]
            total_hits = response['hits']['total']['value']
            
            return {'hits': hits, 'total_hits': total_hits}
    
    # semantic_search is now inherited from base SparseStore class
    # The base class automatically detects chunk_for_semantic_search attribute
    # and uses the appropriate chunking behavior

    # get_db() method removed - use store methods instead

