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_id = meta.get('table_identifier', meta.get('table_number', 'unknown'))
            title = meta.get('table_title', '')
            section = meta.get('section', '')
            source = f"📊 {doc_id} - {table_id}"
            if title:
                source += f": {title}"
            if section:
                source += f" ({section})"
            sources.append(source)
        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(sources)  # Don't use set() to preserve order

def preprocess_query(question):
    import re
    
    question_lower = question.lower()
    
    # Enhanced table detection with appendix
    table_patterns = [
        r'табли[цу]\w*\s+([а-яa-z0-9\.]+)(?:\s+(?:из\s+)?приложени[яеий]\s+(\d+|[а-я]))?',
        r'табли[цу]\w*\s+(?:№|номер)?\s*([а-яa-z0-9\.]+)',
    ]
    
    doc_match = re.search(r'(гост|нп|му)[^\s]*\s*[рp№-]*\s*([0-9\.-]+)', question_lower)
    
    enhanced_query = question
    
    for pattern in table_patterns:
        table_match = re.search(pattern, question_lower)
        if table_match:
            table_num = table_match.group(1).upper()
            enhanced_query += f" таблица номер {table_num}"
            
            # Add appendix context if mentioned
            if len(table_match.groups()) > 1 and table_match.group(2):
                appendix_num = table_match.group(2).upper()
                enhanced_query += f" приложение {appendix_num}"
            break
    
    if doc_match:
        doc_id = f"{doc_match.group(1).upper()} {doc_match.group(2)}"
        enhanced_query += f" документ {doc_id}"
        # Add variations for better matching
        enhanced_query += f" {doc_match.group(1).upper()}Р {doc_match.group(2)}"
    
    return enhanced_query

def answer_question(question, query_engine, reranker):
    try:
        log_message(f"\n{'='*70}")
        log_message(f"Query: {question}")
        log_message(f"{'='*70}")
        
        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"\n📥 INITIAL RETRIEVAL: {len(retrieved)} nodes")
        
        # Detailed logging
        doc_ids = {}
        for n in retrieved:
            doc_id = n.metadata.get('document_id', 'unknown')
            if doc_id not in doc_ids:
                doc_ids[doc_id] = {'tables': [], 'text': 0, 'images': 0}
            
            if n.metadata.get('type') == 'table':
                table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', ''))
                doc_ids[doc_id]['tables'].append(table_id)
            elif n.metadata.get('type') == 'image':
                doc_ids[doc_id]['images'] += 1
            else:
                doc_ids[doc_id]['text'] += 1
        
        for doc_id, counts in doc_ids.items():
            log_message(f"  📄 {doc_id}:")
            if counts['tables']:
                log_message(f"    Tables: {', '.join(set(counts['tables']))}")
            if counts['text']:
                log_message(f"    Text chunks: {counts['text']}")
            if counts['images']:
                log_message(f"    Images: {counts['images']}")
        
        reranked = rerank_nodes(question, retrieved, reranker, top_k=25)
        log_message(f"\n🔄 AFTER RERANKING: {len(reranked)} nodes")
        
        # Detailed reranking results
        doc_ids_reranked = {}
        for n in reranked:
            doc_id = n.metadata.get('document_id', 'unknown')
            if doc_id not in doc_ids_reranked:
                doc_ids_reranked[doc_id] = {'tables': [], 'text': 0, 'images': 0}
            
            if n.metadata.get('type') == 'table':
                table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', ''))
                doc_ids_reranked[doc_id]['tables'].append(table_id)
            elif n.metadata.get('type') == 'image':
                doc_ids_reranked[doc_id]['images'] += 1
            else:
                doc_ids_reranked[doc_id]['text'] += 1
        
        for doc_id, counts in doc_ids_reranked.items():
            log_message(f"  📄 {doc_id}:")
            if counts['tables']:
                log_message(f"    Tables: {', '.join(set(counts['tables']))}")
            if counts['text']:
                log_message(f"    Text chunks: {counts['text']}")
            if counts['images']:
                log_message(f"    Images: {counts['images']}")
        context = "\n\n" + ("="*70 + "\n\n").join(doc_ids_reranked)
        
        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