import gradio as gr
import asyncio
import os
from typing import List
from agent_core import create_agent_graph
from local_knowledge_base import LocalKnowledgeBase
from langchain_core.messages import HumanMessage

class AgenticRAGInterface:
    def __init__(self, agent_graph, knowledge_base):
        self.agent = agent_graph
        self.kb = knowledge_base
        self.conversation_history = []
    
    def process_query(self, query: str, use_rag: bool):
        """Обработка запроса пользователя"""
        
        # Добавляем сообщение в историю
        self.conversation_history.append({"role": "user", "content": query})
        
        # Если нужен RAG, ищем в базе знаний
        context = ""
        if use_rag:
            search_results = self.kb.search(query)
            if search_results:
                context = "\n".join([f"[{i+1}] {res['text'][:200]}..." 
                                    for i, res in enumerate(search_results)])
        
        # Подготавливаем состояние
        initial_state = {
            "messages": [HumanMessage(content=query)],
            "knowledge_base": context,
            "needs_search": use_rag,
            "current_step": "start",
            "final_answer": None
        }
        
        # Запускаем агента
        try:
            result = self.agent.invoke(initial_state)
            answer = result.get("final_answer", "Не удалось получить ответ")
            
            # Добавляем ответ в историю
            self.conversation_history.append({"role": "assistant", "content": answer})
            
            # Форматируем историю для отображения
            history_text = self._format_history()
            
            return answer, history_text
            
        except Exception as e:
            error_msg = f"Ошибка: {str(e)}"
            return error_msg, self._format_history()
    
    def _format_history(self):
        """Форматирование истории разговора"""
        formatted = []
        for msg in self.conversation_history[-10:]:  # Последние 10 сообщений
            role = "👤 Пользователь" if msg["role"] == "user" else "🤖 Ассистент"
            formatted.append(f"{role}: {msg['content']}")
        return "\n\n".join(formatted)
    
    def clear_history(self):
        """Очистка истории"""
        self.conversation_history = []
        return "История очищена", ""

    def index_documents(self, project_name: str, folder_path: str):
        """Индексация документов из указанной папки для проекта"""
        if not project_name:
            return "❌ Укажите имя проекта", ""
        
        if not folder_path:
            return "❌ Укажите путь к папке", ""
        
        if not os.path.exists(folder_path):
            return f"❌ Папка не существует: {folder_path}", ""
        
        if not os.path.isdir(folder_path):
            return f"❌ Указанный путь не является папкой: {folder_path}", ""
        
        try:
            # Переключаемся на проект (создаст новую коллекцию)
            self.kb.set_project(project_name)
            
            # Индексируем документы
            chunks = self.kb.load_documents(folder_path)
            return f"✅ Проект '{project_name}': индексировано чанков: {len(chunks)}", folder_path
        except Exception as e:
            return f"❌ Ошибка при индексации: {str(e)}", ""

    def switch_project(self, project_name: str):
        """Переключиться на другой проект"""
        if not project_name:
            return "❌ Укажите имя проекта", ""
        
        try:
            self.kb.set_project(project_name)
            return f"✅ Переключено на проект: {project_name}", project_name
        except Exception as e:
            return f"❌ Ошибка: {str(e)}", ""

# Инициализируем базу
kb = LocalKnowledgeBase(project_name="default")

# Создаем интерфейс
interface = AgenticRAGInterface(create_agent_graph(kb), kb)

# Создаем Gradio интерфейс
def create_gradio_interface():
    with gr.Blocks(title="Локальный Agentic RAG", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🤖 Локальный Agentic RAG Система
        Полностью автономный AI-агент с базой знаний. Работает без интернета!
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                query_input = gr.Textbox(
                    label="Ваш запрос",
                    placeholder="Задайте вопрос или дайте задание...",
                    lines=3
                )
                
                rag_toggle = gr.Checkbox(
                    label="Использовать базу знаний (RAG)",
                    value=True
                )
                
                submit_btn = gr.Button("Отправить", variant="primary")
                clear_btn = gr.Button("Очистить историю")
                
            with gr.Column(scale=3):
                answer_output = gr.Textbox(
                    label="Ответ агента",
                    lines=8,
                    interactive=False
                )
                
                history_output = gr.Textbox(
                    label="История разговора",
                    lines=12,
                    interactive=False
                )
        
        # Примеры запросов
        gr.Examples(
            examples=[
                ["Объясни концепцию machine learning простыми словами", True],
                ["Посчитай: (15 * 4) + (120 / 3)", False],
                ["Найди информацию о нейронных сетях в базе знаний", True],
                ["Спланируй изучение Python на месяц", False]
            ],
            inputs=[query_input, rag_toggle],
            label="Примеры запросов"
        )
        
        # Секция управления проектами
        gr.Markdown("---")
        gr.Markdown("### 📂 Управление проектами")
        
        with gr.Row():
            with gr.Column(scale=2):
                project_input = gr.Textbox(
                    label="Имя проекта",
                    placeholder="my_cpp_project",
                    lines=1
                )
                switch_btn = gr.Button("🔄 Переключить", variant="secondary")
            with gr.Column(scale=1):
                project_status = gr.Textbox(
                    label="Текущий проект",
                    lines=1,
                    interactive=False
                )
        
        # Секция индексации документов
        gr.Markdown("---")
        gr.Markdown("### 📥 Индексация документов")
        
        with gr.Row():
            with gr.Column(scale=3):
                folder_input = gr.Textbox(
                    label="Путь к папке с файлами",
                    placeholder="./documents или /home/user/my_project",
                    lines=1
                )
            with gr.Column(scale=1):
                index_btn = gr.Button("📥 Индексировать", variant="primary")
        
        status_output = gr.Textbox(
            label="Статус индексации",
            lines=2,
            interactive=False
        )
        
        # Обработчики событий
        submit_btn.click(
            fn=interface.process_query,
            inputs=[query_input, rag_toggle],
            outputs=[answer_output, history_output]
        )
        
        clear_btn.click(
            fn=interface.clear_history,
            inputs=[],
            outputs=[answer_output, history_output]
        )
        
        query_input.submit(
            fn=interface.process_query,
            inputs=[query_input, rag_toggle],
            outputs=[answer_output, history_output]
        )
        
        # Обработчик переключения проекта
        switch_btn.click(
            fn=interface.switch_project,
            inputs=[project_input],
            outputs=[project_status, project_input]
        )
        
        # Обработчик индексации
        index_btn.click(
            fn=interface.index_documents,
            inputs=[project_input, folder_input],
            outputs=[status_output, folder_input]
        )
        
        return demo