import gradio as gr
from huggingface_hub import hf_hub_download
import faiss
import pandas as pd
import os
import json
from llama_index.core import Document, VectorStoreIndex, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer, ResponseMode
from llama_index.core.prompts import PromptTemplate
import time
import sys
from config import *

REPO_ID = "MrSimple01/AIEXP_RAG_FILES"
faiss_index_filename = "cleaned_faiss_index.index"
chunks_filename = "processed_chunks.csv"
download_dir = "rag_files"
table_data_dir = "Табличные данные_JSON"
HF_TOKEN = os.getenv('HF_TOKEN')
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')

CUSTOM_PROMPT_NEW = """
Вы являетесь высокоспециализированным Ассистентом для анализа документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы на основе анализа нормативной документации (НД). Все ваши ответы должны основываться исключительно на предоставленном контексте без использования внешних знаний или предположений.

КРИТИЧЕСКИ ВАЖНО: ВСЕ ОТВЕТЫ ДОЛЖНЫ БЫТЬ ТОЛЬКО НА РУССКОМ ЯЗЫКЕ! НИКОГДА НЕ ОТВЕЧАЙТЕ НА АНГЛИЙСКОМ!

История чата:
{chat_history}


ИНСТРУКЦИИ ПО ОБРАБОТКЕ КОНТЕКСТА:

1. АНАЛИЗ ТАБЛИЧНЫХ ДАННЫХ:
   - Если в контексте есть информация начинающаяся с "Таблица", внимательно изучите её содержимое
   - Извлекайте данные из строк с заголовками и данными таблицы
   - Указывайте номер и название таблицы при ответе
   - Структурируйте ответ на основе табличных данных

2. ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:
Проанализируйте запрос пользователя и определите тип задачи:

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

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

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

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

5. УТОЧНЯЮЩИЕ ВОПРОСЫ (ключевые слова: "что это значит", "что означает", "объясните", "расскажите подробнее"):
   - Используйте историю чата для понимания контекста
   - Если вопрос относится к предыдущему обсуждению, опирайтесь на него
   - Предоставьте подробное объяснение на основе НД
   - Если контекст неясен, попросите уточнения

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

1. ОБЯЗАТЕЛЬНОЕ УКАЗАНИЕ ИСТОЧНИКОВ:
   - Для каждого ответа указывайте: "Согласно [Название документа], раздел [X], пункт [X.X]: [Ваш ответ]"
   - В конце ответа добавляйте: "Подробнее об этом можно узнать в документе [Название документа], раздел [X]."
   - При отсутствии точного раздела: "Согласно документу [Название]: [Ваш ответ]"

2. СТРОГОЕ СЛЕДОВАНИЕ КОНТЕКСТУ:
   - Если информация не найдена: "Информация по вашему запросу не была найдена в нормативной документации."
   - НЕ используйте английский язык ни при каких обстоятельствах
   - Используйте историю чата для понимания контекста вопросов

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

4. СТРУКТУРИРОВАНИЕ ОТВЕТОВ:
   - Основной ответ на русском языке
   - Указание источника
   - Дополнительная информация о документе

Контекст: {context_str}

Вопрос: {query_str}

Ответ (ТОЛЬКО НА РУССКОМ ЯЗЫКЕ):
"""


query_engine = None
chunks_df = None
chat_history = []

def log_message(message):
    print(message, flush=True)
    sys.stdout.flush()

def table_to_document(table_json):
    document_id = table_json.get("document_id") or table_json.get("document", "unknown")
    
    metadata = {
        "document_id": document_id,
        "section": table_json.get("section", ""),
        "table_number": table_json.get("table_number", ""),
        "table_title": table_json.get("table_title", ""),
    }
    
    description = table_json.get("table_description", "")
    headers = table_json.get("headers", [])
    
    table_text = f"ТАБЛИЦА: {table_json.get('table_number', '')} - {table_json.get('table_title', '')}\n"
    table_text += f"ДОКУМЕНТ: {document_id}\n"
    table_text += f"РАЗДЕЛ: {table_json.get('section', '')}\n"
    
    if description:
        table_text += f"ОПИСАНИЕ: {description}\n"
    
    if headers:
        table_text += f"ЗАГОЛОВКИ ТАБЛИЦЫ: {' | '.join(headers)}\n"
    
    data = table_json.get("data", [])
    if data:
        table_text += "ДАННЫЕ ТАБЛИЦЫ:\n"
        for i, row in enumerate(data):
            if isinstance(row, dict):
                row_str = " | ".join([f"{k}: {v}" for k,v in row.items()])
                table_text += f"Строка {i+1}: {row_str}\n"
    
    return Document(text=table_text, metadata=metadata)

def download_table_data():
    log_message("📥 Загрузка табличных данных...")
    
    from huggingface_hub import list_repo_files
    
    table_files = []
    try:
        files = list_repo_files(repo_id=REPO_ID, repo_type="dataset", token=HF_TOKEN)
        for file in files:
            if file.startswith(table_data_dir) and file.endswith('.json'):
                table_files.append(file)
        
        log_message(f"📊 Найдено {len(table_files)} JSON файлов с таблицами")
        
        table_documents = []
        for file_path in table_files:
            try:
                log_message(f"🔄 Обработка файла: {file_path}")
                
                local_path = hf_hub_download(
                    repo_id=REPO_ID,
                    filename=file_path,
                    local_dir=download_dir,
                    repo_type="dataset",
                    token=HF_TOKEN
                )
                
                with open(local_path, 'r', encoding='utf-8') as f:
                    table_data = json.load(f)
                    
                    log_message(f"📋 Структура JSON: {list(table_data.keys()) if isinstance(table_data, dict) else 'Список'}")
                    
                    if isinstance(table_data, dict):
                        if 'sheets' in table_data:
                            for sheet in table_data['sheets']:
                                doc = table_to_document(sheet)
                                table_documents.append(doc)
                                log_message(f"✅ Создан документ из таблицы: {sheet.get('table_number', 'unknown')}")
                        else:
                            doc = table_to_document(table_data)
                            table_documents.append(doc)
                            log_message(f"✅ Создан документ из JSON объекта")
                    elif isinstance(table_data, list):
                        for table_json in table_data:
                            doc = table_to_document(table_json)
                            table_documents.append(doc)
                            log_message(f"✅ Создан документ из элемента списка")
                        
            except Exception as e:
                log_message(f"❌ Ошибка обработки файла {file_path}: {str(e)}")
                continue
        
        log_message(f"✅ Создано {len(table_documents)} документов из таблиц")
        return table_documents
        
    except Exception as e:
        log_message(f"❌ Ошибка загрузки табличных данных: {str(e)}")
        return []


def improve_query_with_history(question, chat_history_list):
    try:
        log_message("🔄 Улучшение запроса с учетом истории...")
        
        if not chat_history_list:
            log_message("📝 История чата пуста, используем оригинальный запрос")
            return question
        
        history_context = ""
        for i, (user_msg, bot_msg) in enumerate(chat_history_list[-3:], 1):
            history_context += f"Вопрос {i}: {user_msg}\nОтвет {i}: {bot_msg[:200]}...\n\n"
        
        improvement_prompt = f"""Проанализируй историю диалога и улучши текущий запрос пользователя.

ИСТОРИЯ ДИАЛОГА:
{history_context}

ТЕКУЩИЙ ЗАПРОС: {question}

ПРАВИЛА:
1. Если запрос неполный или ссылается на предыдущий контекст (например: "что это", "о чем это", "объясни это"), дополни его информацией из истории
2. Если запрос самодостаточный, верни его без изменений  
3. Сохраняй ключевые термины и названия документов из истории
4. Отвечай только улучшенным запросом без дополнительных пояснений

УЛУЧШЕННЫЙ ЗАПРОС:"""

        from llama_index.llms.google_genai import GoogleGenAI
        llm = GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
        
        improved_query = llm.complete(improvement_prompt).text.strip()
        log_message(f"✨ Улучшенный запрос: {improved_query}")
        
        return improved_query
        
    except Exception as e:
        log_message(f"❌ Ошибка улучшения запроса: {str(e)}")
        return question


def format_chat_history():
    if not chat_history:
        return "История чата пуста."
    
    history_text = ""
    for i, (user_msg, bot_msg) in enumerate(chat_history[-5:], 1):
        history_text += f"Сообщение {i}:\nПользователь: {user_msg}\nАссистент: {bot_msg}\n\n"
    
    return history_text


def answer_question(question, history):
    global query_engine, chunks_df, chat_history
    
    if query_engine is None:
        return history + [["", "❌ Система не инициализирована"]], ""
    
    try:
        start_time = time.time()
        
        log_message(f"🔍 Получен вопрос: {question}")
        log_message(f"📜 История чата: {len(chat_history)} сообщений")
        
        # Улучшаем запрос с учетом истории
        improved_question = improve_query_with_history(question, chat_history)
        log_message(f"🎯 Обработка улучшенного запроса: {improved_question}")
        
        # Форматируем историю чата для промпта
        chat_history_text = format_chat_history()
        log_message(f"📝 Сформированная история для промпта: {len(chat_history_text)} символов")
        
        log_message("🔎 Поиск релевантных чанков...")
        retrieved_nodes = query_engine.retriever.retrieve(improved_question)
        log_message(f"📊 Найдено {len(retrieved_nodes)} релевантных чанков")
        
        # Логируем найденные чанки
        for i, node in enumerate(retrieved_nodes[:3]):
            log_message(f"📄 Чанк {i+1}: {node.text[:100]}...")
            log_message(f"🏷️ Метаданные: {node.metadata}")
        
        log_message("🤖 Отправка запроса в LLM...")
        
        # Создаем контекст с историей чата
        query_with_context = f"""
История чата:
{chat_history_text}

Текущий вопрос: {question}
"""
        
        response = query_engine.query(query_with_context)
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        bot_response = response.response
        log_message(f"✅ Получен ответ: {bot_response[:100]}...")
        
        # Проверяем, что ответ на русском языке
        if any(english_word in bot_response.lower() for english_word in ['i am sorry', 'i cannot', 'the query', 'this request']):
            log_message("⚠️ Обнаружен ответ на английском языке, форсируем русский ответ")
            
            # Принудительно запрашиваем ответ на русском
            russian_prompt = f"""
            ВАЖНО: Отвечай ТОЛЬКО на русском языке!
            
            Вопрос: {question}
            История: {chat_history_text}
            Контекст: {retrieved_nodes[0].text if retrieved_nodes else 'Нет контекста'}
            
            Если информации недостаточно для ответа, скажи: "Недостаточно информации для ответа на ваш вопрос в предоставленной документации."
            
            Ответ на русском языке:
            """
            
            from llama_index.llms.google_genai import GoogleGenAI
            llm = GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
            bot_response = llm.complete(russian_prompt).text.strip()
            log_message(f"🔄 Исправленный ответ на русском: {bot_response[:100]}...")
        
        # Обновляем историю чата
        chat_history.append((question, bot_response))
        
        if len(chat_history) > 10:
            chat_history = chat_history[-10:]
        
        log_message(f"💾 История чата обновлена. Всего сообщений: {len(chat_history)}")
        
        sources_html = generate_sources_html(retrieved_nodes)
        
        response_with_time = f"{bot_response}\n\n⏱️ Время обработки: {processing_time:.2f} сек"
        
        history.append([question, response_with_time])
        
        return history, sources_html
        
    except Exception as e:
        error_msg = f"❌ Ошибка обработки вопроса: {str(e)}"
        log_message(f"❌ Ошибка: {str(e)}")
        history.append([question, error_msg])
        return history, ""


def initialize_models():
    global query_engine, chunks_df
    
    try:
        log_message("🔄 Инициализация системы...")
        os.makedirs(download_dir, exist_ok=True)
        
        log_message("📥 Загрузка основных файлов...")
        faiss_index_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=faiss_index_filename,
            local_dir=download_dir,
            repo_type="dataset",
            token=HF_TOKEN
        )
        
        chunks_csv_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=chunks_filename,
            local_dir=download_dir,
            repo_type="dataset",
            token=HF_TOKEN
        )
        
        log_message("📚 Загрузка индекса и данных...")
        index_faiss = faiss.read_index(faiss_index_path)
        chunks_df = pd.read_csv(chunks_csv_path)
        log_message(f"📄 Загружено {len(chunks_df)} основных чанков")
        log_message(f"📋 Колонки в chunks_df: {list(chunks_df.columns)}")
        
        table_documents = download_table_data()
        
        log_message("🤖 Настройка моделей...")
        embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
        llm = GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
        
        Settings.embed_model = embed_model
        Settings.llm = llm
        
        text_column = None
        for col in chunks_df.columns:
            if 'text' in col.lower() or 'content' in col.lower() or 'chunk' in col.lower():
                text_column = col
                break
        
        if text_column is None:
            text_column = chunks_df.columns[0]
        
        log_message(f"📝 Используется колонка для текста: {text_column}")
        
        documents = []
        for i, (_, row) in enumerate(chunks_df.iterrows()):
            doc = Document(
                text=str(row[text_column]), 
                metadata={
                    "chunk_id": row.get('chunk_id', i), 
                    "document_id": row.get('document_id', 'unknown')
                }
            )
            documents.append(doc)
        
        documents.extend(table_documents)
        log_message(f"📋 Всего создано {len(documents)} документов ({len(chunks_df)} чанков + {len(table_documents)} таблиц)")
        
        log_message("🔍 Построение векторного индекса...")
        vector_index = VectorStoreIndex.from_documents(documents)
        
        retriever = VectorIndexRetriever(
            index=vector_index, 
            similarity_top_k=20,
            similarity_cutoff=0.7
        )
        
        custom_prompt_template = PromptTemplate(CUSTOM_PROMPT_NEW)
        response_synthesizer = get_response_synthesizer(
            response_mode=ResponseMode.TREE_SUMMARIZE,
            text_qa_template=custom_prompt_template
        )
        
        query_engine = RetrieverQueryEngine(
            retriever=retriever,
            response_synthesizer=response_synthesizer
        )
        
        log_message("✅ Система успешно инициализирована!")
        return True
        
    except Exception as e:
        log_message(f"❌ Ошибка инициализации: {str(e)}")
        return False

def generate_sources_html(nodes):
    html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; max-height: 400px; overflow-y: auto;'>"
    html += "<h3 style='color: #63b3ed; margin-top: 0;'>📚 Источники:</h3>"
    
    unique_docs = {}
    for node in nodes:
        metadata = node.metadata if hasattr(node, 'metadata') else {}
        doc_id = metadata.get('document_id', 'unknown')
        if doc_id not in unique_docs:
            unique_docs[doc_id] = []
        unique_docs[doc_id].append(node)
    
    for doc_id, doc_nodes in unique_docs.items():
        if doc_id == 'unknown' or doc_id == 'Раздел документа':
            continue
            
        file_link = None
        if chunks_df is not None and 'file_link' in chunks_df.columns:
            doc_rows = chunks_df[chunks_df['document_id'] == doc_id]
            if not doc_rows.empty:
                file_link = doc_rows.iloc[0]['file_link']
        
        html += f"<div style='margin-bottom: 15px; padding: 15px; border: 1px solid #4a5568; border-radius: 8px; background-color: #1a202c;'>"
        html += f"<h4 style='margin: 0 0 10px 0; color: #63b3ed;'>📄 {doc_id}</h4>"
        
        if file_link:
            html += f"<a href='{file_link}' target='_blank' style='color: #68d391; text-decoration: none; font-size: 14px; display: inline-block; margin-bottom: 10px;'>🔗 Ссылка на документ</a><br>"
        
        table_nodes = [node for node in doc_nodes if 'table_number' in node.metadata]
        if table_nodes:
            for node in table_nodes[:3]:
                metadata = node.metadata
                table_num = metadata.get('table_number', '')
                table_title = metadata.get('table_title', 'Без названия')
                if table_num and table_title != 'Без названия':
                    html += f"<p style='font-size: 12px; color: #a0aec0; margin: 5px 0;'>📊 {table_num}: {table_title}</p>"
        
        html += "</div>"
    
    html += "</div>"
    return html

def clear_chat():
    global chat_history
    chat_history = []
    log_message("🗑️ История чата очищена")
    return [], ""

def handle_submit(message, history):
    if not message.strip():
        return history, ""
    
    updated_history, sources = answer_question(message, history)
    return updated_history, sources

def create_demo_interface():
    with gr.Blocks(title="AIEXP - AI Expert для нормативной документации", theme=gr.themes.Soft()) as demo:
        
        gr.Markdown("""
        # AIEXP - Artificial Intelligence Expert
        
        ## Инструмент для работы с нормативной документацией
        """)
        
        with gr.Tab("💬 Чат с документами"):
            gr.Markdown("### Задайте вопрос по нормативной документации")
            
            with gr.Row():
                with gr.Column(scale=2):
                    chatbot = gr.Chatbot(
                        label="Диалог с AIEXP",
                        height=500,
                        show_copy_button=True
                    )
                    
                    with gr.Row():
                        msg = gr.Textbox(
                            label="Ваш вопрос",
                            placeholder="Введите вопрос по нормативным документам...",
                            lines=2,
                            scale=4
                        )
                        send_btn = gr.Button("📤 Отправить", variant="primary", scale=1)
                    
                    with gr.Row():
                        clear_btn = gr.Button("🗑️ Очистить чат", variant="secondary")
                    
                    gr.Examples(
                        examples=[
                            "Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?",
                            "Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?",
                            "В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями, не включенными в перечисления?",
                        ],
                        inputs=msg
                    )
                
                with gr.Column(scale=1):
                    sources_output = gr.HTML(
                        label="Источники",
                        value="<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Здесь появятся источники...</div>",
                    )
            
            def submit_message(message, history):
                return handle_submit(message, history)
            
            msg.submit(submit_message, [msg, chatbot], [chatbot, sources_output]).then(
                lambda: "", None, msg
            )
            
            send_btn.click(submit_message, [msg, chatbot], [chatbot, sources_output]).then(
                lambda: "", None, msg
            )
            
            clear_btn.click(clear_chat, outputs=[chatbot, sources_output])
    
    return demo

if __name__ == "__main__":
    log_message("🚀 Запуск AIEXP - AI Expert для нормативной документации")
    
    if initialize_models():
        log_message("🌟 Запуск веб-интерфейса...")
        demo = create_demo_interface()
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            debug=False
        )
    else:
        log_message("❌ Невозможно запустить приложение из-за ошибки инициализации")
        sys.exit(1)