# solution.py

from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langchain.tools import tool
from langchain_ollama import ChatOllama
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command
from rich import print as rprint

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

# ------------------------------------------------------------------
# 2. Создаём LLM (Ollama) и агент с HumanInTheLoopMiddleware
# ------------------------------------------------------------------
llm = ChatOllama(model="llama3")          # используем Ollama

memory = MemorySaver()                    # память для сохранения пауз

agent = create_agent(
    model=llm,
    tools=[get_weather],
    system_prompt="Ты полезный ассистент, помогающий пользователю.",
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={
                "get_weather": True          # все решения: approve, edit, reject
            },
            description_prefix="Подтвердите вызов инструмента",
        ),
    ],
    checkpointer=memory,
)

# ------------------------------------------------------------------
# 3. Функция для запроса решений у пользователя
# ------------------------------------------------------------------
def ask_decisions(action_requests):
    """
    action_requests – список словарей с полями name, args и опционально description.
    Возвращает список решений в том же порядке: approve/reject.
    """
    decisions = []
    for idx, act in enumerate(action_requests, start=1):
        rprint(f"\n--- Подтверждение #{idx} ---")
        rprint(f"[bold]Инструмент:[/bold] {act['name']}")
        rprint(f"[bold]Аргументы:[/bold] {act['args']}")
        if "description" in act:
            rprint(f"[bold]Описание:[/bold] {act['description']}")

        while True:
            choice = input("a = approve, r = reject: ").strip().lower()
            if choice == "a":
                decisions.append({"type": "approve"})
                break
            elif choice == "r":
                msg = input("Сообщение для агента (причина отказа): ")
                decisions.append({"type": "reject", "message": msg})
                break
            else:
                rprint("[red]Неверный ввод. Пожалуйста, введите 'a' или 'r'.[/red]")
    return decisions

# ------------------------------------------------------------------
# 4. Основной цикл взаимодействия с агентом
# ------------------------------------------------------------------
def main():
    thread_id = "сессия-1"
    config = {"configurable": {"thread_id": thread_id}}

    while True:
        user_msg = input("\nВы: ")
        if not user_msg.strip():
            continue

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

        # цикл обработки пауз
        while "__interrupt__" in result:
            interrupt_value = result["__interrupt__"][0].value
            action_requests = interrupt_value.get("action_requests", [])
            if not action_requests:
                rprint("[yellow]Пауза без действий. Продолжаем...[/yellow]")
                break

            decisions = ask_decisions(action_requests)

            # возобновляем выполнение
            result = agent.invoke(
                Command(resume={"decisions": decisions}),
                config=config,
            )

        # вывод финального ответа агента
        if "messages" in result and result["messages"]:
            last_msg = result["messages"][-1]
            rprint(f"\n[bold]Агент:[/bold] {last_msg['content']}")
        else:
            rprint("[red]Ответ не получен.[/red]")

if __name__ == "__main__":
    main()