# solution.py

"""
Пример агента с Human-in-the-Loop через middleware.
Запускается из командной строки и демонстрирует работу интерактивного подтверждения вызова инструмента.
"""

from langgraph.prebuilt import create_react_agent
from langchain.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command, StateGraph
from langgraph.moderation import HumanInTheLoopMiddleware

# ------------------------------------------------------------
# 1. Определяем простой инструмент
# ------------------------------------------------------------
@tool
def get_weather(city: str, date: str) -> str:
    """Возвращает погоду в городе на указанную дату."""
    return f"Погода в {city} на {date}: солнечно 25°C."

# ------------------------------------------------------------
# 2. Создаём агент с HumanInTheLoopMiddleware
# ------------------------------------------------------------
memory = MemorySaver()

agent = create_react_agent(
    tools=[get_weather],
    system_prompt="Ты полезный ассистент, отвечай точно.",
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={"get_weather": True},  # прерываем только при вызове get_weather
            description_prefix="Подтвердите вызов инструмента",
        ),
    ],
    checkpointer=memory,
)

# ------------------------------------------------------------
# 3. Запускаем агент и обрабатываем интерактивный цикл
# ------------------------------------------------------------
def run_agent():
    """
    Вводим сообщение от пользователя, запускаем агента,
    обрабатываем прерывания и продолжаем работу до завершения.
    """
    # Инициализируем конфиг с thread_id
    config = {"configurable": {"thread_id": "session-1"}}

    # Запрашиваем ввод от пользователя
    user_input = input("Введите запрос: ")

    # Первый вызов агента
    result = agent.invoke(
        {"messages": [{"role": "human", "content": user_input}]},
        config=config,
    )

    # Цикл обработки прерываний
    while "__interrupt__" in result:
        interrupt_payload = result["__interrupt__"][0].value
        action_requests = interrupt_payload.get("action_requests", [])

        decisions = []
        for req in action_requests:
            print("\nИнструмент:", req["name"])
            print("Аргументы:", req["args"])

            # Запрашиваем решение пользователя
            decision = input("a=approve, r=reject: ").strip().lower()
            if decision == "r":
                reason = input("Причина отказа: ").strip()
                decisions.append({"type": "reject", "message": reason})
            else:
                decisions.append({"type": "approve"})

        # Возобновляем работу агента
        result = agent.invoke(Command(resume={"decisions": decisions}), config=config)

    # Выводим финальный ответ агента
    final_message = result["messages"][-1]["content"]
    print("\nОтвет агента:", final_message)


if __name__ == "__main__":
    run_agent()