# solution.py

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

from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
from rich.console import Console

console = Console()

# ---------- 1. Настройка модели и инструментов ----------
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0,
)

def echo_tool(args: Dict[str, Any]) -> str:
    """Простейший инструмент для демонстрации."""
    return f"Эхо: {args.get('message', '')}"

tools = [echo_tool]
tool_executor = ToolExecutor(tools)


# ---------- 2. Создание графа агента ----------
def create_agent() -> StateGraph:
    """
    Возвращает готовый граф агента с памятью и возможностью
    прерывания перед вызовом инструмента.
    """
    # Состояние: хранит сообщения
    def add_message(state: Dict[str, Any], message: Dict[str, str]) -> Dict[str, Any]:
        state["messages"].append(message)
        return state

    graph = StateGraph(add_message)

    # Входной узел – пользовательское сообщение
    graph.set_entry_point("user")

    # Узлы для обработки сообщений и вызова инструментов
    graph.add_node("assistant", lambda state: {"messages": [state["messages"][-1]]})
    graph.add_node("tools", tool_executor)

    # Переходы
    graph.set_conditional_edges(
        "assistant",
        lambda _: "tools" if any(m.get("tool_calls") for m in _["messages"]) else "final",
    )
    graph.add_edge("tools", "assistant")
    graph.add_edge("assistant", "final")

    # Конфигурация прерывания
    graph.set_interrupt_before(["tools"])

    return graph


agent = create_agent()
memory = MemorySaver()
agent.compile(checkpointer=memory)


# ---------- 3. Функция взаимодействия ----------
def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]) -> None:
    """
    Запускает потоковый ответ агента с возможностью подтверждения
    вызова инструмента пользователем.
    """
    # Инициализируем сообщение пользователя в состоянии
    if user_input is not None:
        agent.send(user_input, config=config)

    for chunk_type, chunk_data in agent.stream(
        {}, config=config, stream_mode=["messages", "updates"]
    ):
        state = agent.get_state(config)

        # 3.1 Вывод токенов LLM
        if chunk_type == "messages":
            token = chunk_data["content"] or ""
            console.print(token, end="", style="cyan")
            continue

        # 3.2 Обработка обновлений (вызовы инструментов)
        if chunk_type == "updates":
            for update in chunk_data:
                console.print(f"\n[bold magenta]Обновление:{update}")
            continue

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


# ---------- 4. Основной цикл ----------
def main() -> None:
    thread_id = "thread-1"
    config = {"configurable": {"thread_id": thread_id}}

    console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")

    while True:
        user_input = input("\nВы: ")
        if user_input.lower() == "exit":
            break
        ask_and_run(
            {"messages": [{"role": "human", "content": user_input}]},
            config,
        )
    console.print("[bold green]До свидания!")


if __name__ == "__main__":
    main()