# solution.py
# -*- coding: utf-8 -*-
"""
Пример агента с Human‑in‑the‑loop, реализованный в LangGraph.
"""

from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command
from langchain.tools import tool
from langchain_core.messages import ToolMessage


# 1. Модель LLM (пример с OpenAI‑совместимым API)
llm = ChatOpenAI(
    model="gpt-4o-mini",  # замените на нужную модель
    temperature=0.7,
)

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

# 3. Хранилище памяти для чекпоинтера
memory = MemorySaver()

# 4. Создание агента с interrupt_before=["tools"]
agent = create_react_agent(
    tools=[get_weather],
    model=llm,
    system_prompt="Ты полезный ассистент, отвечай точно.",
    checkpointer=memory,
    interrupt_before=["tools"],  # включаем human‑in‑the‑loop
)

# 5. Основная логика взаимодействия
def main() -> None:
    config = {"configurable": {"thread_id": "session-1"}}
    while True:
        user_input = input("Вы: ")
        if user_input.lower() in {"выход", "exit", "quit"}:
            print("Завершение работы.")
            break
        # первый вызов агента
        result = agent.invoke(
            {"messages": [{"role": "human", "content": user_input}]},
            config=config,
        )
        # цикл обработки прерываний
        while "__interrupt__" in result:
            last_msg = result["messages"][-1]
            tool_call = last_msg.tool_calls[0]  # один вызов инструмента
            name = tool_call["name"]
            args = tool_call["args"]
            print(f"\nИнструмент: {name}")
            print(f"Аргументы: {args}")
            choice = input("a=approve, r=reject: ").strip().lower()
            if choice == "r":
                reason = input("Причина отказа: ")
                # передаем ToolMessage как результат вызова инструмента
                result = agent.invoke(
                    {"messages": [ToolMessage(content=reason, tool_call_id=tool_call["id"])]},
                    config=config,
                )
            else:
                # approve – просто возобновляем без сообщения
                result = agent.invoke(Command(resume=None), config=config)
        # вывод финального ответа агента
        final_answer = result["messages"][-1].content
        print("\nОтвет агента:\n", final_answer)

if __name__ == "__main__":
    main()