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 preprocess_query(question):
    import re
    
    question_lower = question.lower()
    
    table_match = re.search(r'табли[цу]\w*\s+([а-яa-z0-9\.]+)', question_lower)
    doc_match = re.search(r'(гост|нп|му)[^\s]*\s*[рp№-]*\s*([0-9\.-]+)', question_lower)
    
    enhanced_query = question
    
    if table_match:
        table_num = table_match.group(1).upper()
        enhanced_query += f" таблица номер {table_num}"
    
    if doc_match:
        doc_id = f"{doc_match.group(1).upper()} {doc_match.group(2)}"
        enhanced_query += f" документ {doc_id}"
    
    return enhanced_query

def answer_question(question, query_engine, reranker):
    try:
        log_message(f"Query: {question}")
        
        enhanced_query = preprocess_query(question)
        if enhanced_query != question:
            log_message(f"Enhanced query: {enhanced_query}")
        
        retrieved = query_engine.retriever.retrieve(enhanced_query)
        log_message(f"Retrieved {len(retrieved)} nodes")
        
        doc_ids = [n.metadata.get('document_id', 'unknown') for n in retrieved]
        table_nums = [n.metadata.get('table_number', '') for n in retrieved if n.metadata.get('type') == 'table']
        log_message(f"Retrieved from documents: {set(doc_ids)}")
        if table_nums:
            log_message(f"Retrieved tables: {set(table_nums)}")
        
        reranked = rerank_nodes(question, retrieved, reranker, top_k=25)
        log_message(f"Reranked to {len(reranked)} nodes")
        
        doc_ids_reranked = [n.metadata.get('document_id', 'unknown') for n in reranked]
        table_nums_reranked = [n.metadata.get('table_number', '') for n in reranked if n.metadata.get('type') == 'table']
        log_message(f"After reranking - documents: {set(doc_ids_reranked)}")
        if table_nums_reranked:
            log_message(f"After reranking - tables: {set(table_nums_reranked)}")
        
        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_num = meta.get('table_number', 'unknown')
                title = meta.get('table_title', '')
                source_label = f"[ТАБЛИЦА {table_num} - {doc_id}]"
                if title:
                    source_label += f" {title}"
            elif doc_type == 'image':
                img_num = meta.get('image_number', 'unknown')
                source_label = f"[РИСУНОК {img_num} - {doc_id}]"
            else:
                section = meta.get('section_id', '')
                source_label = f"[{doc_id} - {section}]"
            
            context_parts.append(f"{source_label}\n{n.text}")
        
        context = "\n\n" + ("="*70 + "\n\n").join(context_parts)
        
        prompt = f"""Ты эксперт по технической документации. 

КОНТЕКСТ:
{context}

ВОПРОС: {question}

ИНСТРУКЦИИ:
1. Используй ТОЛЬКО контекст выше
2. Если спрашивают содержание таблицы - ОБЯЗАТЕЛЬНО приведи ВСЕ данные из таблицы
3. Укажи источник: документ и номер таблицы
4. Если таблица разбита на части - объедини информацию
5. Если информации нет - четко скажи об этом

ОТВЕТ:"""
        
        response = query_engine.query(prompt)
        sources = format_sources(reranked)
        
        return response.response, 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=20, min_score=0.3):
    """Rerank nodes with diversity - MORE LENIENT"""
    if not nodes:
        return []
    
    # Score all nodes
    pairs = [[query, n.text] for n in nodes]
    scores = reranker.predict(pairs)
    
    # Sort by score
    scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
    
    # More lenient threshold
    filtered = [(n, s) for n, s in scored if s >= min_score]
    
    if not filtered:
        # Fallback: take top 50% if nothing passes threshold
        cutoff = max(scores) * 0.5
        filtered = [(n, s) for n, s in scored if s >= cutoff][:top_k]
    
    # Log top scores for debugging
    log_message(f"Top 5 reranking scores: {[f'{s:.3f}' for _, s in scored[:5]]}")
    
    # Diversity selection - but prioritize tables if query mentions them
    selected = []
    seen_docs = set()
    table_nodes = []
    other_nodes = []
    
    for node, score in filtered:
        if node.metadata.get('type') == 'table':
            table_nodes.append((node, score))
        else:
            other_nodes.append((node, score))
    
    # If query mentions "таблица", prioritize table nodes
    if 'таблиц' in query.lower():
        combined = table_nodes + other_nodes
    else:
        combined = filtered
    
    for node, score in combined[:top_k]:
        if len(selected) >= top_k:
            break
        selected.append(node)
        seen_docs.add(node.metadata.get('document_id', 'unknown'))
    
    log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(seen_docs)} docs)")
    
    return selected