# solution.py
# -*- coding: utf-8 -*-
"""
Агент с HumanInTheLoopMiddleware: при каждом вызове инструмента
агент останавливается, пользователь подтверждает (approve/reject),
после чего выполнение возобновляется через Command.
"""

import json
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

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

# 2. LLM
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0,
)

# 3. Память
memory = MemorySaver()

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

# 5. Сбор решений от пользователя
def get_user_decisions(action_requests: list[dict], review_configs: list[dict]) -> list[dict]:
    decisions = []
    for action, review_cfg in zip(action_requests, review_configs):
        name = action.get("name", "unknown")
        args = action.get("args", {})
        description = action.get("description", "")
        allowed = review_cfg.get("allowed_decisions", ["approve", "reject"])

        print(f"\n--- Подтверждение ---")
        print(f"Инструмент: {name}")
        print(f"Аргументы: {json.dumps(args, ensure_ascii=False)}")
        if description:
            print(f"Описание: {description}")
        print(f"Разрешённые решения: {', '.join(allowed)}")

        while True:
            choice = input("a=approve, r=reject: ").strip().lower()
            if choice in ("a", "approve"):
                decisions.append({"type": "approve"})
                break
            elif choice in ("r", "reject"):
                msg = input("Причина отказа: ").strip()
                decisions.append({
                    "type": "reject",
                    "message": msg or "Запрос отклонён пользователем",
                })
                break
            else:
                print("Неверный ввод. Введите 'a' или 'r'.")
    return decisions

# 6. Основной цикл
def main() -> None:
    config = {"configurable": {"thread_id": "session-1"}}
    print("Привет! Введите запрос или 'выход' для завершения.")

    while True:
        user_input = input("\nВы: ").strip()
        if user_input.lower() in {"выход", "exit", "quit"}:
            print("Завершение работы.")
            break

        result = agent.invoke(
            {"messages": [{"role": "human", "content": user_input}]},
            config=config,
        )

        while "__interrupt__" in result:
            interrupt_value = result["__interrupt__"][0].value
            action_requests = interrupt_value.get("action_requests", [])
            review_configs = interrupt_value.get("review_configs", [])

            decisions = get_user_decisions(action_requests, review_configs)

            result = agent.invoke(
                Command(resume={"decisions": decisions}),
                config=config,
            )

        print(f"\nАгент: {result['messages'][-1].content}")

if __name__ == "__main__":
    main()