from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer, ResponseMode
from llama_index.core.prompts import PromptTemplate
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from my_logging import log_message
from config import CUSTOM_PROMPT, PROMPT_SIMPLE_POISK

def create_vector_index(documents):
    log_message("Строю векторный индекс")
    return VectorStoreIndex.from_documents(documents)

def create_query_engine(vector_index):
    bm25_retriever = BM25Retriever.from_defaults(
        docstore=vector_index.docstore, 
        similarity_top_k=80
    )
    
    vector_retriever = VectorIndexRetriever(
        index=vector_index, 
        similarity_top_k=80,
        similarity_cutoff=0.45
    )
    
    hybrid_retriever = QueryFusionRetriever(
        [vector_retriever, bm25_retriever],
        similarity_top_k=100,
        num_queries=1
    )
    
    custom_prompt_template = PromptTemplate(PROMPT_SIMPLE_POISK)
    response_synthesizer = get_response_synthesizer(
        response_mode=ResponseMode.TREE_SUMMARIZE,  
        text_qa_template=custom_prompt_template
    )
    
    query_engine = RetrieverQueryEngine(
        retriever=hybrid_retriever,
        response_synthesizer=response_synthesizer
    )
    
    return query_engine

    
def rerank_nodes(query, nodes, reranker, top_k=40, min_score_threshold=0.35, diversity_penalty=0.15):
    if not nodes or not reranker:
        return nodes[:top_k]
    
    pairs = [[query, node.text] for node in nodes]
    scores = reranker.predict(pairs)
    scored_nodes = list(zip(nodes, scores))
    scored_nodes.sort(key=lambda x: x[1], reverse=True)
    
    if min_score_threshold:
        scored_nodes = [(node, score) for node, score in scored_nodes 
                      if score >= min_score_threshold]
    
    if not scored_nodes:
        scored_nodes = list(zip(nodes, scores))
        scored_nodes.sort(key=lambda x: x[1], reverse=True)
        scored_nodes = scored_nodes[:top_k]
    
    selected = []
    seen_docs = {}
    
    for node, score in scored_nodes:
        if len(selected) >= top_k:
            break
        
        meta = node.metadata if hasattr(node, 'metadata') else {}
        doc_id = meta.get('document_id', 'unknown')
        node_type = meta.get('type', 'text')
        table_num = meta.get('table_number', '')
        
        key = f"{doc_id}_{table_num}" if node_type == 'table' else f"{doc_id}_{meta.get('section_id', '')}"
        
        if key in seen_docs:
            penalty = diversity_penalty * 0.2 if node_type == 'table' else diversity_penalty
            adjusted_score = score * (1 - penalty)
        else:
            adjusted_score = score
            seen_docs[key] = 1
        
        if not selected or adjusted_score >= selected[0][1] * 0.4:
            selected.append((node, score))
    
    return [node for node, score in selected]