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):
    try:
        bm25_retriever = BM25Retriever.from_defaults(
            docstore=vector_index.docstore, 
            similarity_top_k=30 
        )
        
        vector_retriever = VectorIndexRetriever(
            index=vector_index, 
            similarity_top_k=30,  
            similarity_cutoff=0.65
        )
        
        hybrid_retriever = QueryFusionRetriever(
            [vector_retriever, bm25_retriever],
            similarity_top_k=40,  
            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
        )
        
        log_message("Query engine успешно создан")
        return query_engine
        
    except Exception as e:
        log_message(f"Ошибка создания query engine: {str(e)}")
        raise

def rerank_nodes(query, nodes, reranker, top_k=25, min_score_threshold=0.45, diversity_penalty=0.2):
    """Rerank with better handling of specific technical queries"""
    if not nodes or not reranker:
        return nodes[:top_k]
    
    try:
        log_message(f"Переранжирую {len(nodes)} узлов для запроса: {query[:50]}...")
        
        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)
        
        # Lower threshold for technical queries
        if min_score_threshold is not None:
            scored_nodes = [(node, score) for node, score in scored_nodes 
                          if score >= min_score_threshold]
            log_message(f"После фильтрации (порог {min_score_threshold}): {len(scored_nodes)} узлов")
        
        if not scored_nodes:
            log_message("⚠️ Нет узлов после фильтрации, снижаю порог до 0.3")
            scored_nodes = list(zip(nodes, scores))
            scored_nodes.sort(key=lambda x: x[1], reverse=True)
            min_score_threshold = max(0.3, scored_nodes[0][1] * 0.5)
            scored_nodes = [(node, score) for node, score in scored_nodes 
                          if score >= min_score_threshold]
        
        selected_nodes = []
        selected_docs = {}  # Track count per document
        selected_tables = set()
        
        for node, score in scored_nodes:
            if len(selected_nodes) >= top_k:
                break
            
            metadata = node.metadata if hasattr(node, 'metadata') else {}
            doc_id = metadata.get('document_id', 'unknown')
            node_type = metadata.get('type', 'text')
            
            # Track table uniqueness
            if node_type == 'table':
                table_id = metadata.get('full_table_id', '')
                if table_id in selected_tables:
                    continue  # Skip duplicate table chunks
                selected_tables.add(table_id)
            
            # Apply lighter diversity penalty
            penalty = 0
            doc_count = selected_docs.get(doc_id, 0)
            if doc_count > 0:
                penalty = min(diversity_penalty * doc_count, 0.5)
            
            adjusted_score = score * (1 - penalty)
            
            # Accept if competitive
            if not selected_nodes or adjusted_score >= selected_nodes[0][1] * 0.5:
                selected_nodes.append((node, score))
                selected_docs[doc_id] = doc_count + 1
        
        log_message(f"✓ Выбрано {len(selected_nodes)} узлов")
        log_message(f"  Уникальных документов: {len(selected_docs)}")
        log_message(f"  Уникальных таблиц: {len(selected_tables)}")
        
        if selected_nodes:
            log_message(f"  Score: {selected_nodes[0][1]:.3f} → {selected_nodes[-1][1]:.3f}")
        
        return [node for node, score in selected_nodes]
        
    except Exception as e:
        log_message(f"❌ Ошибка переранжировки: {str(e)}")
        return nodes[:top_k]