from llama_index.llms.google_genai import GoogleGenAI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from sentence_transformers import CrossEncoder
from my_logging import log_message

def get_llm_model(api_key, model_name="gemini-2.0-flash"):
    """Get LLM model"""
    return GoogleGenAI(model=model_name, api_key=api_key)

def get_embedding_model(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"):
    """Get embedding model"""
    return HuggingFaceEmbedding(model_name=model_name)

def get_reranker_model(model_name='cross-encoder/ms-marco-MiniLM-L-12-v2'):
    """Get reranker model"""
    return CrossEncoder(model_name)


def format_sources(nodes):
    """Format retrieved sources for display"""
    sources = []
    for node in nodes:
        meta = node.metadata
        doc_type = meta.get('type', 'text')
        doc_id = meta.get('document_id', 'unknown')
        
        if doc_type == 'table':
            table_num = meta.get('table_number', 'unknown')
            title = meta.get('table_title', '')
            sources.append(f"📊 {doc_id} - Таблица {table_num}: {title}")
        elif doc_type == 'image':
            img_num = meta.get('image_number', 'unknown')
            sources.append(f"🖼️ {doc_id} - Рисунок {img_num}")
        else:
            section = meta.get('section_id', '')
            sources.append(f"📄 {doc_id} - Раздел {section}")
    
    return "\n".join(set(sources))

def answer_question(question, query_engine, reranker):
    try:
        log_message(f"\n{'='*70}")
        log_message(f"QUERY: {question}")

        # Retrieve and rerank nodes
        retrieved = query_engine.retriever.retrieve(question)
        log_message(f"\nRETRIEVED: {len(retrieved)} nodes")
        reranked = rerank_nodes(question, retrieved, reranker, top_k=25, min_score=0.3)
        log_message(f"\nRERANKED: {len(reranked)} nodes")

        # Build context for prompt
        context_parts = []
        for n in reranked:
            meta = n.metadata
            doc_id = meta.get('document_id', 'unknown')
            doc_type = meta.get('type', 'text')
            if doc_type == 'table':
                table_id = meta.get('table_identifier', meta.get('table_number', 'unknown'))
                title = meta.get('table_title', '')
                source_label = f"[{doc_id} - Таблица {table_id}]"
                if title:
                    source_label += f" {title}"
            else:
                source_label = f"[{doc_id}]"
            context_parts.append(f"{source_label}\n{n.text}")  # Use FULL text, not [:500]

        context = "\n\n" + ("="*50 + "\n\n").join(context_parts)

        # Use CUSTOM_PROMPT from config
        from config import CUSTOM_PROMPT
        prompt = CUSTOM_PROMPT.format(context_str=context, query_str=question)
        log_message(f"\nPROMPT LENGTH: {len(prompt)} chars\n")
        
        # CRITICAL FIX: Call LLM directly instead of query_engine.query()
        from llama_index.core import Settings
        response = Settings.llm.complete(prompt)

        sources = format_sources(reranked)
        
        # Log retrieved chunks
        log_message(f"\n{'='*70}")
        log_message("RETRIEVED CHUNKS:")
        for i, node in enumerate(reranked, 1):
            log_message(f"\n--- Chunk {i} ---")
            log_message(f"Document: {node.metadata.get('document_id', 'unknown')}")
            log_message(f"Type: {node.metadata.get('type', 'unknown')}")
            if node.metadata.get('type') == 'table':
                log_message(f"Table: {node.metadata.get('table_identifier', 'unknown')}")
            log_message(f"Text preview: {node.text[:500]}...")
        
        return response.text, sources

    except Exception as e:
        log_message(f"Error: {e}")
        import traceback
        log_message(traceback.format_exc())
        return f"Ошибка: {e}", ""
    

def rerank_nodes(query, nodes, reranker, top_k=25, min_score=0.3):
    """Simple and effective reranking: sort by score and filter by threshold."""
    if not nodes or not reranker:
        return nodes[:top_k]

    pairs = [[query, n.text] for n in nodes]
    scores = reranker.predict(pairs)
    scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
    filtered = [n for n, s in scored if s >= min_score]

    # Return top_k filtered nodes, or fallback to top_k overall
    return filtered[:top_k] if filtered else [n for n, _ in scored[:top_k]]