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

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 better deduplication"""
    log_message("Creating query engine...")

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

    class DeduplicatedQueryEngine(RetrieverQueryEngine):
        def retrieve(self, query):
            nodes = hybrid_retriever.retrieve(query)
            log_message(f"Hybrid retrieval returned: {len(nodes)} nodes")
            
            # Better deduplication using longer text snippet
            seen_hashes = set()
            unique_nodes = []
            doc_type_counts = {'text': 0, 'table': 0, 'image': 0}
            
            for node in nodes:
                # Use first 500 chars for dedup hash
                text_hash = hash(node.text[:500])
                
                if text_hash not in seen_hashes:
                    seen_hashes.add(text_hash)
                    unique_nodes.append(node)
                    
                    # Count by type
                    node_type = node.metadata.get('type', 'text')
                    doc_type_counts[node_type] = doc_type_counts.get(node_type, 0) + 1
            
            log_message(f"After dedup: {len(unique_nodes)} unique nodes")
            log_message(f"Types: text={doc_type_counts.get('text', 0)}, "
                       f"table={doc_type_counts.get('table', 0)}, "
                       f"image={doc_type_counts.get('image', 0)}")
            
            return unique_nodes[:50]

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

    log_message("✓ Query engine created")
    return query_engine