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

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

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

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

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

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

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

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

5. ДОПОЛНИТЕЛЬНЫЕ РЕКОМЕНДАЦИИ:
   - При множественных релевантных источниках - укажите все
   - Выделяйте критически важные требования
   - Указывайте альтернативные процедуры, если они предусмотрены НД

Контекст: {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):
    metadata = {
        "document_id": table_json["document_id"],
        "section": table_json["section"],
        "table_number": table_json["table_number"],
        "table_title": table_json["table_title"],
    }
    
    description = table_json["table_description"]
    headers = " | ".join(table_json["headers"])
    
    rows = []
    for row in table_json["data"]:
        row_str = " | ".join([f"{k}: {v}" for k,v in row.items()])
        rows.append(row_str)
    
    table_text = f"Таблица {table_json['table_number']} - {table_json['table_title']}\n"
    table_text += f"Описание: {description}\n"
    table_text += f"Заголовки: {headers}\n"
    table_text += "\n".join(rows)
    
    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:
                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)
                    
                    if isinstance(table_data, list):
                        for table_json in table_data:
                            doc = table_to_document(table_json)
                            table_documents.append(doc)
                    else:
                        doc = table_to_document(table_data)
                        table_documents.append(doc)
                        
            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 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 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)} основных чанков")
        
        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("📝 Создание документов из чанков...")
        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 answer_question(question, history):
    global query_engine, chunks_df, chat_history
    
    if query_engine is None:
        return history + [["", "❌ Система не инициализирована"]], ""
    
    try:
        start_time = time.time()
        
        chat_history_text = format_chat_history()
        
        query_with_history = question
        response = query_engine.query(query_with_history)
        retrieved_nodes = query_engine.retriever.retrieve(query_with_history)
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        bot_response = response.response
        chat_history.append((question, bot_response))
        
        if len(chat_history) > 10:
            chat_history = chat_history[-10:]
        
        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)}"
        history.append([question, error_msg])
        return history, ""

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():
        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>"
        
        for node in doc_nodes[:3]:
            metadata = node.metadata if hasattr(node, 'metadata') else {}
            if 'table_number' in metadata:
                html += f"<p style='font-size: 12px; color: #a0aec0; margin: 5px 0;'>📊 Таблица {metadata['table_number']}: {metadata.get('table_title', 'Без названия')}</p>"
        
        html += "</div>"
    
    html += "</div>"
    return html

def clear_chat():
    global chat_history
    chat_history = []
    return [], ""

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 user_message(message, history):
                return "", history + [[message, None]]
            
            def bot_message(history):
                if history and history[-1][1] is None:
                    user_msg = history[-1][0]
                    updated_history, sources = answer_question(user_msg, history[:-1])
                    return updated_history, sources
                return history, ""
            
            msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot_message, chatbot, [chatbot, sources_output]
            )
            
            send_btn.click(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot_message, chatbot, [chatbot, sources_output]
            )
            
            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)