from qdrant_client import QdrantClient


class HybridSearcher:
    DENSE_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
    SPARSE_MODEL = "prithivida/Splade_PP_en_v1"
    def __init__(self, collection_name):
        self.collection_name = collection_name
        # initialize Qdrant client
        self.qdrant_client = QdrantClient("http://localhost:6333")
        self.qdrant_client.set_model(self.DENSE_MODEL)
        # comment this line to use dense vectors only
        self.qdrant_client.set_sparse_model(self.SPARSE_MODEL)
        
    def search(self, text: str, city: str):
        city_of_interest = city

        # Define a filter for cities
        city_filter = models.Filter(
            must=[
                models.FieldCondition(
                    key="city", 
                    match=models.MatchValue(value=city_of_interest)
                )
            ]
        )
    
        search_result = self.qdrant_client.query(
            collection_name=self.collection_name,
            query_text=text,
            query_filter=city_filter,
            limit=5
        )
        # `search_result` contains found vector ids with similarity scores 
        # along with the stored payload
        
        # Select and return metadata
        metadata = [hit.metadata for hit in search_result]
        return metadata