import gradio as gr
import asyncio
from typing import List

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 "История очищена", ""

# Создаем интерфейс
interface = AgenticRAGInterface(agent_graph, 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="Примеры запросов"
        )
        
        # Обработчики событий
        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]
        )
        
        return demo

# Запускаем интерфейс
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False  # Не создавать публичную ссылку
    )