import os
from typing import Any, Dict, Iterable

# Библиотеки для вывода и работы с агентом
from rich.console import Console
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langgraph.prebuilt import create_react_agent

# Инициализация консоли для красивого вывода
console = Console()

# Настройка модели LLM (потребуется переменная окружения OPENAI_API_KEY)
llm = ChatOpenAI(temperature=0.7)

# Простой инструмент – эхо‑ответ
from langchain_core.tools import tool

@tool
def echo_tool(text: str) -> str:
    """Возвращает введённый текст."""
    return text

# Создаём агент с памятью и возможностью паузы перед вызовом инструмента
memory = MemorySaver()
agent = create_react_agent(
    model=llm,
    tools=[echo_tool],
    prompt="You are a helpful assistant. Use the echo tool when needed.",
    checkpointer=memory,          # обязательно для памяти и паузы
    interrupt_before=["tools"],   # остановка перед каждым вызовом инструмента
)

# Конфигурация разговора (thread_id)
config = {"configurable": {"thread_id": "chat-1"}}

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

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

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

        # Обнаруживаем паузу перед инструментом
        if "__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, cfg)
            else:
                console.print("[bold red]Отменено пользователем.")
                break

# Основной цикл чата
if __name__ == "__main__":
    console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")
    while True:
        user_text = input("\nВы: ")
        if user_text.lower() == "exit":
            console.print("[bold magenta]До свидания!")
            break
        # Формируем сообщение пользователя
        msg = {"messages": [{"role": "human", "content": user_text}]}
        ask_and_run(msg, config)