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()
    
    # Extract document ID and normalize
    doc_match = re.search(r'(гост|нп|му)\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
    
    enhanced_query = question
    
    if doc_match:
        doc_type = doc_match.group(1).upper()
        doc_num = doc_match.group(2)
        
        # Add normalized versions
        enhanced_query += f" {doc_type} Р {doc_num}"
    
    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_stats = {}
        for n in retrieved:
            doc_id = n.metadata.get('document_id', 'unknown')
            doc_type = n.metadata.get('type', 'text')
            
            if doc_id not in doc_stats:
                doc_stats[doc_id] = {'tables': set(), 'text': 0, 'images': 0}
            
            if doc_type == 'table':
                table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
                doc_stats[doc_id]['tables'].add(table_id)
            elif doc_type == 'image':
                doc_stats[doc_id]['images'] += 1
            else:
                doc_stats[doc_id]['text'] += 1
        
        for doc_id in sorted(doc_stats.keys()):
            stats = doc_stats[doc_id]
            parts = []
            if stats['tables']:
                parts.append(f"tables={list(stats['tables'])[:5]}")
            if stats['text']:
                parts.append(f"text={stats['text']}")
            if stats['images']:
                parts.append(f"images={stats['images']}")
            log_message(f"  {doc_id}: {', '.join(parts)}")

        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)
        from config import CUSTOM_PROMPT
        prompt = f"""
Вы являетесь высокоспециализированным Ассистентом для анализа нормативных документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы исключительно на основе предоставленного контекста из нормативной документации.

ПРАВИЛА АНАЛИЗА ЗАПРОСА:

1. ПРЯМЫЕ ВОПРОСЫ БЕЗ ДОКУМЕНТАЛЬНОГО КОНТЕКСТА:
   Если пользователь задает вопрос типа "В каких случаях могут быть признаны протоколы испытаний?" без предоставления дополнительных документов, найдите соответствующую информацию в доступном контексте и предоставьте полный ответ с указанием источников.

2. ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:

   а) ПОИСК И ОТВЕТ НА ВОПРОС (ключевые слова: "в каких случаях", "когда", "кто", "что", "как", "почему"):
      - Найдите релевантную информацию в контексте
      - Предоставьте развернутый ответ
      - Обязательно укажите конкретные документы и разделы
      - Процитируйте ключевые положения

   б) КРАТКОЕ САММАРИ (ключевые слова: "кратко", "суммировать", "резюме", "основные моменты"):
      - Предоставьте структурированное резюме
      - Выделите ключевые требования
      - Используйте нумерованный список

   в) ПОИСК ДОКУМЕНТА И ПУНКТА (ключевые слова: "найти", "где", "какой документ", "в каком разделе"):
      - Укажите конкретный документ и структурное расположение
      - Предоставьте точные номера разделов/пунктов

   г) ПРОВЕРКА КОРРЕКТНОСТИ (ключевые слова: "правильно ли", "соответствует ли", "проверить"):
      - Четко укажите: "СООТВЕТСТВУЕТ" или "НЕ СООТВЕТСТВУЕТ"
      - Перечислите конкретные требования

   д) ПЛАН ДЕЙСТВИЙ (ключевые слова: "план", "алгоритм", "пошагово"):
      - Создайте пронумерованный план
      - Укажите ссылки на соответствующие пункты НД

ПРАВИЛА ФОРМИРОВАНИЯ ОТВЕТОВ:

Работай исключительно с информацией из предоставленного контекста. Запрещено использовать:
- Общие знания
- Информацию из интернета
- Данные из предыдущих диалогов
- Собственные предположения

1. СТРУКТУРА ОТВЕТА:
   - Начинайте с прямого ответа на вопрос
   - Затем указывайте нормативные основания
   - Завершайте ссылками на конкретные документы и разделы

2. РАБОТА С КОНТЕКСТОМ:
   - Если информация найдена в контексте - предоставьте полный ответ
   - Если информация не найдена: "Информация по вашему запросу не найдена в доступной нормативной документации"
   - Не делайте предположений за пределами контекста
   - Не используйте общие знания

3. ТЕРМИНОЛОГИЯ И ЦИТИРОВАНИЕ:
   - Сохраняйте официальную терминологию НД
   - Цитируйте точные формулировки ключевых требований
   - При множественных источниках - укажите все релевантные

4. ФОРМАТИРОВАНИЕ:
   - Для перечислений: используйте нумерованные списки
   - Выделяйте критически важные требования
   - Структурируйте ответ логически

# КАК РАБОТАТЬ С ЗАПРОСОМ

**Шаг 1:** Определи, что именно ищет пользователь (термин, требование, процедура, условие)

**Шаг 2:** Найди релевантную информацию в контексте

**Шаг 3:** Сформируй ответ:
   - Если нашел: укажи документ и пункт, процитируй нужную часть
   - Если не нашел: четко сообщи об отсутствии информации

**Шаг 4:** При наличии нескольких источников:
   - Представь их последовательно с указанием источника каждого
   - Если источников много (>4) — сначала дай их список, потом цитаты   

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

ВОПРОС: {question}
"""
        
        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