from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer
from my_logging import log_message

SIMPLE_PROMPT = """Вы - эксперт по нормативной документации.

Контекст:
{context_str}

Вопрос: {query_str}

Инструкция:
1. Отвечайте ТОЛЬКО на основе предоставленного контекста
2. Цитируйте конкретные источники (документ, раздел, таблицу)
3. Если информации недостаточно, четко укажите это
4. Будьте точны и конкретны

Ответ:"""

def create_vector_index(documents):
    """Create vector index from documents"""
    log_message(f"Building vector index from {len(documents)} documents...")
    index = VectorStoreIndex.from_documents(documents)
    log_message("✓ Index created")
    return index

def keyword_filter_nodes(query, nodes, min_keyword_matches=1):
    """Return nodes that contain at least one keyword from the query."""
    keywords = [w.lower() for w in query.split() if len(w) > 2]
    filtered = []
    for node in nodes:
        text = node.text.lower()
        if any(k in text for k in keywords):
            filtered.append(node)
    return filtered

def create_query_engine(vector_index):
    """Create hybrid retrieval engine with deduplication"""
    log_message("Creating query engine...")

    vector_retriever = VectorIndexRetriever(
        index=vector_index,
        similarity_top_k=50  # Reduced from 50
    )
    bm25_retriever = BM25Retriever.from_defaults(
        docstore=vector_index.docstore,
        similarity_top_k=50  # Reduced from 50
    )
    hybrid_retriever = QueryFusionRetriever(
        [vector_retriever, bm25_retriever],
        similarity_top_k=60,  # Reduced from 60
        num_queries=1
    )

    class DeduplicatedQueryEngine(RetrieverQueryEngine):
        def retrieve(self, query):
            nodes = hybrid_retriever.retrieve(query)
            
            # CRITICAL: Deduplicate by text content hash
            seen_hashes = set()
            unique_nodes = []
            
            for node in nodes:
                text_hash = hash(node.text[:200])
                
                if text_hash not in seen_hashes:
                    seen_hashes.add(text_hash)
                    unique_nodes.append(node)
            
            log_message(f"Retrieved: {len(nodes)} → Unique: {len(unique_nodes)}")
            return unique_nodes[:50]  # Return top 50 unique

    response_synthesizer = get_response_synthesizer()
    
    query_engine = DeduplicatedQueryEngine(
        retriever=hybrid_retriever,
        response_synthesizer=response_synthesizer
    )

    log_message("✓ Query engine created (with deduplication)")
    return query_engine