# agent.py

import os
from typing import Any, Dict, Iterable, Tuple

from langchain_ollama import ChatOllama
from langchain_qdrant import QdrantStore
from langchain.embeddings.nomic import NomicEmbeddings
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_chat_agent
from langchain.tools import tool
from rich.console import Console

# ------------------------------------------------------------------
# Конфигурация
# ------------------------------------------------------------------

CONSOLE = Console()

LLM_MODEL = "llama3"
EMBEDDINGS_MODEL = "nomic-embed-text"

QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_COLLECTION = "knowledge_base"

# ------------------------------------------------------------------
# Векторный хранилище
# ------------------------------------------------------------------

embeddings = NomicEmbeddings(model=EMBEDDINGS_MODEL)
vector_store = QdrantStore(
    url=QDRANT_URL,
    collection_name=QDRANT_COLLECTION,
    embeddings=embeddings,
)

# ------------------------------------------------------------------
# Инструменты RAG
# ------------------------------------------------------------------


@tool("search_knowledge_base")
def search_knowledge_base(query: str) -> str:
    """
    Поиск в базе знаний по запросу.
    Возвращает найденный текст (или сообщение об отсутствии результата).
    """
    results = vector_store.similarity_search_with_score(query, k=3)
    if not results:
        return "Ничего не найдено."
    # Собираем топ‑k результатов
    texts = "\n\n".join([f"Score: {score:.2f}\n{doc.page_content}" for doc, score in results])
    return f"Найденные документы:\n{texts}"


@tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str) -> str:
    """
    Добавляет новый текст в базу знаний.
    Возвращает подтверждение добавления.
    """
    vector_store.add_texts([content])
    return "Текст успешно добавлен в базу знаний."


# ------------------------------------------------------------------
# Создание агента
# ------------------------------------------------------------------

llm = ChatOllama(model=LLM_MODEL)

memory = MemorySaver()

def create_agent() -> StateGraph:
    """
    Возвращает готовый граф агента с памятью и паузой перед вызовом инструментов.
    """
    # Базовый агент от LangChain
    base_agent = create_chat_agent(
        llm=llm,
        tools=[search_knowledge_base, add_to_knowledge_base],
        system_prompt="Ты полезный ассистент. При необходимости используй инструменты.",
    )

    # Создаём граф с паузой перед инструментами
    graph = StateGraph()
    graph.add_node("agent", base_agent)
    graph.set_entry_point("agent")
    graph.add_edge("agent", END)

    # Включаем interrupt_before для паузы
    graph.interrupt_before = ["tools"]

    return graph

AGENT_GRAPH = create_agent()


# ------------------------------------------------------------------
# Функция взаимодействия
# ------------------------------------------------------------------


def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]) -> None:
    """
    Обрабатывает ввод пользователя и выводит потоковые ответы.
    При паузе спрашивает подтверждение перед вызовом инструмента.
    """
    for chunk in AGENT_GRAPH.stream(
        user_input,
        config=config,
        stream_mode=["messages", "updates"],
    ):
        state = AGENT_GRAPH.get_state(config)
        chunk_type, chunk_data = chunk

        # Печать токенов
        if chunk_type == "messages":
            CONSOLE.print(chunk_data["content"], end="")

        # Печать вызова инструмента
        if chunk_type == "updates" and "tool_calls" in chunk_data:
            for call in chunk_data["tool_calls"]:
                name = call["name"]
                args = call.get("arguments", {})
                CONSOLE.print(f"\n[bold cyan]Инструмент:{name}({args})[/bold cyan]")

        # Обработка паузы
        if "__interrupt__" in chunk_data and state.next == ("tools",):
            # Получаем последний вызов инструмента
            last_msg = state.values["messages"][-1]
            tool_call = last_msg.tool_calls[0]
            name, args = tool_call["name"], tool_call.get("arguments", {})
            CONSOLE.print(f"\n[bold magenta]Агент хочет вызвать инструмент {name}({args})[/bold magenta]")
            answer = input("Разрешить? (Y/n): ")
            if answer.lower().strip() in ("y", ""):
                # Возобновляем без новых сообщений
                ask_and_run(None, config)
            else:
                CONSOLE.print("[red]Отменено[/red]")
                break


# ------------------------------------------------------------------
# Основной цикл чата
# ------------------------------------------------------------------


def main() -> None:
    thread_id = "thread-1"
    config = {"configurable": {"thread_id": thread_id}}

    while True:
        user_input = input("\nВы: ")
        if user_input.lower().strip() == "exit":
            break

        ask_and_run(
            {"messages": [{"role": "human", "content": user_input}]},
            config,
        )


if __name__ == "__main__":
    main()