from typing import TypedDict, List, Optional, Annotated
import operator
from langchain_community.llms import Ollama
from langchain_community.vectorstores import Qdrant
from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_core.messages import AIMessage, SystemMessage
from langchain.tools import tool
# from langchain.agents import AgentExecutor, create_react_agent
import numpy as np

# Определяем состояние агента
class AgentState(TypedDict):
    messages: Annotated[List, add_messages]  # История сообщений
    knowledge_base: Optional[str]  # Релевантные документы
    current_step: str  # Текущий шаг выполнения
    needs_search: bool  # Нужен ли поиск в RAG
    final_answer: Optional[str]  # Финальный ответ

# Инициализация LLM через Ollama
llm = Ollama(
    model="qwen2.5-coder-16k:14b",
    temperature=0.1,  # Низкая температура для консистентности
    num_predict=1024,  # Максимальная длина ответа
)

# Создаем инструменты для агента
@tool
def search_knowledge_base(query: str) -> str:
    """Поиск информации в локальной базе знаний"""
    # Здесь будет подключение к Qdrant
    return "Найденная информация из базы знаний"

@tool
def calculate(expression: str) -> str:
    """Выполнение математических вычислений"""
    try:
        result = eval(expression)
        return f"Результат: {result}"
    except:
        return "Ошибка в выражении"

@tool
def web_search(query: str) -> str:
    """Поиск в интернете (если нужно)"""
    # Можно подключить локальный поиск через DuckDuckGo
    return "Результаты поиска из интернета"

# Создаем граф агента
def create_agent_graph():
    workflow = StateGraph(AgentState)
    
    # Узел: анализ запроса
    def analyze_query(state: AgentState):
        messages = state["messages"]
        last_message = messages[-1].content if messages else ""
        
        # Простой анализ: проверяем, нужен ли поиск
        search_keywords = ["информация", "документ", "найди", "ищи", "база знаний"]
        needs_search = any(keyword in last_message.lower() for keyword in search_keywords)
        
        return {
            "needs_search": needs_search,
            "current_step": "analyzing_query"
        }
    
    # Узел: поиск в RAG
    def rag_search(state: AgentState):
        if not state["needs_search"]:
            return {"knowledge_base": None, "current_step": "generating_answer"}
        
        # Получаем последний запрос
        query = state["messages"][-1].content
        
        # Здесь должен быть реальный поиск в Qdrant
        # Пока заглушка
        results = ["Документ 1: Информация о...", "Документ 2: Данные по..."]
        
        return {
            "knowledge_base": "\n".join(results),
            "current_step": "generating_answer"
        }
    
    # Узел: генерация ответа
    def generate_answer(state: AgentState):
        messages = state["messages"]
        knowledge = state.get("knowledge_base", "")
        
        # Формируем промпт с контекстом
        prompt = f"""Ты — интеллектуальный ассистент. Используй следующую информацию если она релевантна:
        
        Контекст из базы знаний:
        {knowledge}
        
        История разговора:
        {messages[-5:] if len(messages) > 5 else messages}
        
        Текущий запрос: {messages[-1].content if messages else ''}
        
        Ответь максимально полезно и точно:"""
        
        response = llm.invoke(prompt)
        
        return {
            "final_answer": response,
            "current_step": "completed",
            "messages": messages + [AIMessage(content=response)]
        }
    
    # Добавляем узлы в граф
    workflow.add_node("analyze", analyze_query)
    workflow.add_node("search", rag_search)
    workflow.add_node("generate", generate_answer)
    
    # Определяем edges (переходы)
    workflow.set_entry_point("analyze")
    workflow.add_edge("analyze", "search")
    workflow.add_edge("search", "generate")
    workflow.add_edge("generate", END)
    
    return workflow.compile()