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}")
    
        
        # Detect listing queries - need MORE chunks
        is_listing_query = any(phrase in question.lower() 
                              for phrase in ['какие таблиц', 'список', 'перечисл', 'все таблиц'])
        
        retrieved = query_engine.retriever.retrieve(question)
        log_message(f"\nRETRIEVED: {len(retrieved)} nodes")
        
        # Log retrieved docs
        doc_stats = {}
        for n in retrieved:
            doc_id = n.metadata.get('document_id', 'unknown')
            doc_group = n.metadata.get('document_group', doc_id)
            
            if doc_group not in doc_stats:
                doc_stats[doc_group] = {'tables': set(), 'text': 0}
            
            if n.metadata.get('type') == 'table':
                table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
                doc_stats[doc_group]['tables'].add(table_id)
            else:
                doc_stats[doc_group]['text'] += 1
        
        for doc_id in sorted(doc_stats.keys()):
            stats = doc_stats[doc_id]
            log_message(f"  {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
            if stats['tables']:
                log_message(f"    Tables: {sorted(stats['tables'])}")
        
        # Adjust reranking based on query type
        if is_listing_query:
            reranked = rerank_nodes(question, retrieved, reranker, top_k=50, min_score=0.2)
        else:
            reranked = rerank_nodes(question, retrieved, reranker, top_k=25, min_score=0.3)
        
        log_message(f"\nRERANKED: {len(reranked)} nodes")
        
        # Log reranked
        doc_stats_reranked = {}
        for n in reranked:
            doc_group = n.metadata.get('document_group', n.metadata.get('document_id', 'unknown'))
            
            if doc_group not in doc_stats_reranked:
                doc_stats_reranked[doc_group] = {'tables': set(), 'text': 0}
            
            if n.metadata.get('type') == 'table':
                table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
                doc_stats_reranked[doc_group]['tables'].add(table_id)
            else:
                doc_stats_reranked[doc_group]['text'] += 1
        
        for doc_id in sorted(doc_stats_reranked.keys()):
            stats = doc_stats_reranked[doc_id]
            log_message(f"  {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
            if stats['tables']:
                log_message(f"    Tables: {sorted(stats['tables'])}")
        
        # Build context
        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[:500]}")  # Limit context per chunk
        
        context = "\n\n" + ("="*50 + "\n\n").join(context_parts)
        
        # Adjust prompt for listing queries
        if is_listing_query:
            prompt = f"""Контекст содержит информацию о таблицах из документов.

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

ВОПРОС: {question}

ИНСТРУКЦИИ:
1. Перечисли ВСЕ таблицы, найденные в контексте для запрошенного документа
2. Укажи номер таблицы и название (если есть)
3. Если таблиц нет - скажи прямо

ОТВЕТ (список таблиц):"""
        else:
            prompt = f"""Ты эксперт по технической документации.

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

ВОПРОС: {question}

ИНСТРУКЦИИ:
1. Отвечай ТОЛЬКО на основе контекста
2. Укажи источник (документ, таблицу)
3. Если нужно показать содержимое таблицы - покажи ВСЕ данные
4. Если информации нет - скажи прямо

ОТВЕТ:"""
        
        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=25, min_score=0.3):
    """Rerank with document grouping awareness"""
    if not nodes:
        return []
    
    pairs = [[query, n.text] for n in nodes]
    scores = reranker.predict(pairs)
    
    scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
    
    log_message(f"Top 10 reranking scores: {[f'{s:.3f}' for _, s in scored[:10]]}")
    
    # More lenient filtering
    filtered = [(n, s) for n, s in scored if s >= min_score]
    
    if not filtered:
        cutoff = max(scores) * 0.4
        filtered = [(n, s) for n, s in scored if s >= cutoff][:top_k]
    
    # Group by document for diversity
    doc_groups = {}
    for node, score in filtered:
        doc_group = node.metadata.get('document_group', node.metadata.get('document_id', 'unknown'))
        if doc_group not in doc_groups:
            doc_groups[doc_group] = []
        doc_groups[doc_group].append((node, score))
    
    # Take top chunks from each document group
    selected = []
    group_limits = max(3, top_k // max(1, len(doc_groups)))
    
    for doc_group in doc_groups:
        selected.extend([n for n, s in doc_groups[doc_group][:group_limits]])
    
    # Fill remaining slots with highest scores
    if len(selected) < top_k:
        remaining = [n for n, s in filtered if n not in selected]
        selected.extend(remaining[:top_k - len(selected)])
    
    log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(doc_groups)} doc groups)")
    
    return selected[:top_k]