import gradio as gr
import os
import sys
import logging
from config import *
from documents_prep import DocumentsPreparation
import index_retriever
from chat_handler import ChatHandler

REPO_ID = "MrSimple01/AIEXP_RAG_FILES"
HF_TOKEN = os.getenv('HF_TOKEN')

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

doc_prep = None
chat_handler = None

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

def initialize_system():
    global doc_prep, chat_handler
    
    try:
        log_message("Запуск инициализации системы AIEXP")
        
        doc_prep = DocumentsPreparation(REPO_ID, HF_TOKEN)
        
        log_message("Подготовка документов")
        all_documents = doc_prep.prepare_all_documents()
        
        if not all_documents:
            log_message("Не удалось загрузить документы")
            return False
        
        log_message("Инициализация моделей и индекса")
        if not index_retriever.initialize_models(all_documents):
            log_message("Не удалось инициализировать модели")
            return False
        
        chat_handler = ChatHandler(None)
        
        log_message("Система успешно инициализирована")
        return True
        
    except Exception as e:
        log_message(f"Ошибка инициализации системы: {str(e)}")
        return False

def handle_question(question):
    if chat_handler is None:
        return "Система не инициализирована", ""
    
    try:
        answer = index_retriever.query(question)
        sources = get_sources_for_question(question)
        
        # chat_handler.add_to_history(question, answer)
        
        return answer, sources
    except Exception as e:
        error_msg = f"Ошибка обработки вопроса: {str(e)}"
        log_message(error_msg)
        return error_msg, ""

def get_sources_for_question(question):
    try:
        nodes = index_retriever.retrieve_nodes(question)
        if not nodes:
            return "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Источники не найдены</div>"
        
        sources_html = "<div style='background-color: #2d3748; color: white; padding: 15px; border-radius: 10px;'>"
        sources_html += "<h3 style='color: #4fd1c7; margin-top: 0;'>📚 Источники:</h3>"
        
        for i, node in enumerate(nodes[:5], 1):
            source_text = node.text[:200] + "..." if len(node.text) > 200 else node.text
            sources_html += f"<div style='margin: 10px 0; padding: 10px; background-color: #4a5568; border-radius: 5px;'>"
            sources_html += f"<strong>Источник {i}:</strong><br>"
            sources_html += f"<small>{source_text}</small>"
            sources_html += "</div>"
        
        sources_html += "</div>"
        return sources_html
        
    except Exception as e:
        log_message(f"Ошибка получения источников: {str(e)}")
        return "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Ошибка загрузки источников</div>"

def handle_model_switch(model_name):
    try:
        return index_retriever.switch_model(model_name)
    except Exception as e:
        error_msg = f"Ошибка переключения модели: {str(e)}"
        log_message(error_msg)
        return f"❌ {error_msg}"

def get_current_model_status():
    try:
        if not index_retriever.is_initialized():
            return "Система не инициализирована"
        return f"Текущая модель: {index_retriever.get_current_model()}"
    except Exception as e:
        return "Ошибка получения статуса модели"

def get_chat_history_html():
    if chat_handler is None:
        return "История недоступна"
    return chat_handler.get_history_html()

def clear_chat_history():
    if chat_handler is not None:
        chat_handler.clear_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):
                    model_dropdown = gr.Dropdown(
                        choices=list(AVAILABLE_MODELS.keys()),
                        value=DEFAULT_MODEL,
                        label="🤖 Выберите языковую модель",
                        info="Выберите модель для генерации ответов"
                    )
                with gr.Column(scale=1):
                    switch_btn = gr.Button("🔄 Переключить модель", variant="secondary")
                    model_status = gr.Textbox(
                        value=get_current_model_status(),
                        label="Статус модели",
                        interactive=False
                    )
            
            with gr.Row():
                with gr.Column(scale=3):
                    question_input = gr.Textbox(
                        label="Ваш вопрос к базе знаний",
                        placeholder="Введите вопрос по нормативным документам...",
                        lines=3
                    )
                    ask_btn = gr.Button("🔍 Найти ответ", variant="primary", size="lg")
                    
                    gr.Examples(
                        examples=[
                            "О чем этот рисунок: ГОСТ Р 50.04.07-2022 Приложение Л. Л.1.5 Рисунок Л.2",
                            "Л.9 Формула в ГОСТ Р 50.04.07 - 2022 что и о чем там?", 
                            "Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?",
                            "Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?",
                            "В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями?",
                        ],
                        inputs=question_input
                    )
            
            with gr.Row():
                with gr.Column(scale=2):
                    answer_output = gr.HTML(
                        label="",
                        value=f"<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Здесь появится ответ на ваш вопрос...<br><small>Текущая модель: {DEFAULT_MODEL}</small></div>",
                    )
                
                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>",
                    )
        
        switch_btn.click(
            fn=handle_model_switch,
            inputs=[model_dropdown],
            outputs=[model_status]
        )
        
        ask_btn.click(
            fn=handle_question,
            inputs=[question_input],
            outputs=[answer_output, sources_output]
        )
        
        question_input.submit(
            fn=handle_question,
            inputs=[question_input],
            outputs=[answer_output, sources_output]
        )
    
    return demo

if __name__ == "__main__":
    log_message("Запуск AIEXP - AI Expert для нормативной документации")
    
    if initialize_system():
        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)