# solution.py

import os
from typing import Any, Dict, Optional

from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, MessagesState, END, START
from langgraph.prebuilt import ToolNode
from rich.console import Console

# --------------------------------------------------------------------------- #
# 1. Консоль для вывода
# --------------------------------------------------------------------------- #

console = Console()

# --------------------------------------------------------------------------- #
# 2. LLM
# --------------------------------------------------------------------------- #

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# --------------------------------------------------------------------------- #
# 3. Инструмент
# --------------------------------------------------------------------------- #

@tool
def echo_tool(text: str) -> str:
    """Возвращает введённый текст без изменений."""
    return f"Эхо: {text}"

tools = [echo_tool]

# --------------------------------------------------------------------------- #
# 4. Граф агента (актуальный подход LangGraph + LangChain v1.0+)
# --------------------------------------------------------------------------- #

llm_with_tools = llm.bind_tools(tools)

def call_model(state: MessagesState) -> Dict[str, Any]:
    """Вызывает LLM с историей сообщений."""
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

def should_continue(state: MessagesState) -> str:
    """Определяет следующий шаг: вызов инструмента или завершение."""
    if state["messages"][-1].tool_calls:
        return "tools"
    return END

graph = StateGraph(MessagesState)
graph.add_node("agent", call_model)
graph.add_node("tools", ToolNode(tools))
graph.add_edge(START, "agent")
graph.add_conditional_edges("agent", should_continue)
graph.add_edge("tools", "agent")

memory = MemorySaver()
agent = graph.compile(
    checkpointer=memory,
    interrupt_before=["tools"],  # пауза перед каждым вызовом инструмента
)

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

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

        # 5.1 Потоковый вывод токенов LLM
        if chunk_type == "messages":
            for msg in chunk_data:
                if isinstance(msg, AIMessage):
                    console.print(msg.content, end="")

        # 5.2 Вывод вызовов инструментов
        elif chunk_type == "updates":
            for upd in chunk_data:
                console.print(f"[bold cyan]Tool call: {upd}")

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

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

if __name__ == "__main__":
    config = {"configurable": {"thread_id": "chat-1"}}
    console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")

    while True:
        user_text = input("\nВы: ")
        if user_text.lower() == "exit":
            console.print("[bold magenta]До свидания!")
            break

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