import gradio as gr
import os
import sys
import logging
from config import *
from documents_prep import DocumentsPreparation
from index_retriever import IndexRetriever
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
index_retriever = None
chat_handler = None

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

def initialize_system():
    global doc_prep, index_retriever, chat_handler
    
    try:
        log_message("Запуск инициализации системы AIEXP")
        
        doc_prep = DocumentsPreparation(REPO_ID, HF_TOKEN)
        index_retriever = IndexRetriever(config=sys.modules[__name__])
        
        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(index_retriever)
        
        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 "Система не инициализирована", ""
    return chat_handler.answer_question(question)

def handle_model_switch(model_name):
    if index_retriever is None:
        return "Система не инициализирована"
    return index_retriever.switch_model(model_name)

def get_current_model_status():
    if index_retriever is None:
        return "Система не инициализирована"
    return f"Текущая модель: {index_retriever.get_current_model()}"

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)