from qdrant_client import QdrantClient
from fastembed import SparseTextEmbedding
from qdrant_client import QdrantClient, models
from reranker import Reranker
from sentence_transformers import SentenceTransformer
from config import DENSE_MODEL, SPARSE_MODEL, QDRANT_URL, QDRANT_API_KEY
import math

class DocSearcherV2:
    
    def __init__(self, collection_name):
        self.collection_name = collection_name
        self.reranker = Reranker()
        self.model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B",device="cpu")
        # self.sparse_model = SparseTextEmbedding(SPARSE_MODEL)  # Disabled - not needed without sparse search
        self.qdrant_client = QdrantClient(QDRANT_URL,api_key=QDRANT_API_KEY,timeout=30)
    
    async def search_semantic(self, text: str, qdrant_limit: int = 20, top_k: int = 3):
        """
        Semantička pretraga sa reranking-om
        
        Args:
            text: Query tekst
            qdrant_limit: Broj chunk-ova iz Qdrant-a (default: 20)
            top_k: Broj najboljih rezultata nakon rerankinga (default: 3)
        """
        queries = [text]
        dense_query = self.model.encode(text).tolist()
        # sparse_query = next(self.sparse_model.query_embed(text))  # Disabled - collection not configured

        # Dense-only search (sparse disabled due to collection config)
        prefetch = [
                models.Prefetch(
                    query=dense_query,
                    using="Qwen/Qwen3-Embedding-0.6B",
                    limit=qdrant_limit
                ),
        ]
            
        search_result = self.qdrant_client.query_points(
            collection_name= "sl-list",
            prefetch=prefetch,
            query=models.FusionQuery(
                fusion=models.Fusion.RRF,
            ),
            with_payload=True,
            limit=qdrant_limit,
        ).points

        # Sačuvaj kompletan hit objekat i tekst za reranking
        # Kreiraj mapu tekst -> hit za brzo mapiranje
        text_to_hit = {}
        texts_for_reranking = []

        for hit in search_result:
            text = hit.payload.get("tekst", "")
            if text:  # Samo ako postoji tekst
                text_to_hit[text] = hit
                texts_for_reranking.append(text)
        
        if not texts_for_reranking:
            return []
            
        # Reranking - ISPRAVLJENO: jedan query za sve dokumente
        reranked_results = self.reranker.compute_logits(
            queries * len(texts_for_reranking),  # Svaki dokument dobija isti query
            texts_for_reranking,
            top_k=top_k  # Vrati samo top_k rezultata
        )
        
        # Kombinuj rezultate: mapiraj rerank skorove sa originalnim hit-ovima
        # reranked_results je lista tuple-ova: [(score, query, document_text), ...]
        # gde je document_text originalni tekst koji je poslat reranker-u
        results_with_scores = []
        for score, query, document_text in reranked_results:
            # Pronađi originalni hit po tekstu
            if document_text in text_to_hit:
                hit = text_to_hit[document_text]
                
                # Sanitizuj score - osiguraj da je validna float vrijednost za JSON
                score_float = float(score)
                if math.isnan(score_float) or math.isinf(score_float):
                    score_float = 0.0  # Default za invalid skorove
                elif score_float < 0:
                    score_float = 0.0
                elif score_float > 1:
                    score_float = 1.0
                
                # Vrati kompletan payload sa skorom
                result = {
                    "score": score_float,
                    "id": str(hit.id),
                    "text": document_text,
                    "payload": hit.payload  # Kompletan payload sa svim podacima
                }
                results_with_scores.append(result)
        
        # Reranker već vraća sortirano, ali osigurajmo da je sortirano
        results_with_scores.sort(key=lambda x: x["score"], reverse=True)
        
        # Vrati top rezultate (reranker već vraća top 10)
        return results_with_scores
    
    async def search_keyword(self, text: str):
        sparse_query = next(self.sparse_model.query_embed(text))

        prefetch = [
                models.Prefetch(
                    query=models.SparseVector(**sparse_query.as_object()),
                    using=SPARSE_MODEL,
                    limit=100
                )
            ]
            
        search_result = self.qdrant_client.query_points(
            collection_name= "sl-list",
            prefetch=prefetch,
            query=models.FusionQuery(
                fusion=models.Fusion.RRF,
            ),
            with_payload=True,
            limit = 100,
        ).points

        data = []

        for hit in search_result:
            data.append(hit.payload["tekst"])
        
        return data