"""
Human-in-the-Loop через middleware

Стек строго по условию задания:
  - langchain.agents.create_agent
  - langchain.agents.middleware.HumanInTheLoopMiddleware
  - langgraph.checkpoint.memory.MemorySaver
  - langgraph.types.Command
  - result["__interrupt__"] для проверки паузы
  - Command(resume={"decisions": [...]}) для возобновления
"""

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

# ── Модель ───────────────────────────────────────────────────────────────────
# ChatOpenAI поддерживает вызов инструментов (tool calling), что необходимо агенту.
# Ключ читается из переменной окружения OPENAI_API_KEY.
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

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

# ── Агент с HumanInTheLoopMiddleware ─────────────────────────────────────────
# checkpointer обязателен: без него пауза не сохраняется
memory = MemorySaver()

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

# ── Основной чат-цикл ────────────────────────────────────────────────────────
def main():
    # thread_id обязателен — привязывает состояние к одной сессии,
    # благодаря чему пауза сохраняется и выполнение можно возобновить
    config = {"configurable": {"thread_id": "сессия-1"}}

    print("Привет! Я ассистент. Введите 'exit' для выхода.")

    while True:
        try:
            user_text = input("\nВы: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nДо свидания!")
            break

        if not user_text:
            continue
        if user_text.lower() == "exit":
            print("До свидания!")
            break

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

        # ── Цикл обработки прерываний ─────────────────────────────────────
        # Пока агент приостановлен для подтверждения — показываем действие,
        # собираем решение пользователя и возобновляем через Command.
        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 = []

            for action in action_requests:
                name        = action.get("name", "")
                args        = action.get("args", {})
                description = action.get("description", "")

                print("\n--- Подтверждение ---")
                print(f"Инструмент: {name}")
                print(f"Аргументы:  {args}")
                if description:
                    print(f"Описание:   {description}")

                # Определяем допустимые решения для этого действия
                allowed = ["approve", "reject"]
                for cfg in review_configs:
                    if cfg.get("name") == name:
                        allowed = cfg.get("allowed_decisions", allowed)
                        break

                # Формируем подсказку для пользователя
                hints = []
                if "approve" in allowed:
                    hints.append("a = approve")
                if "reject"  in allowed:
                    hints.append("r = reject")
                if "edit"    in allowed:
                    hints.append("e = edit")

                while True:
                    choice = input(", ".join(hints) + ": ").strip().lower()

                    if choice in ("a", "approve") and "approve" in allowed:
                        decisions.append({"type": "approve"})
                        break

                    elif choice in ("r", "reject") and "reject" in allowed:
                        message = input(
                            "Сообщение для агента (причина отказа): "
                        ).strip()
                        decisions.append({
                            "type": "reject",
                            "message": message or "Нет причины",
                        })
                        break

                    elif choice in ("e", "edit") and "edit" in allowed:
                        # По желанию: редактирование аргументов перед выполнением
                        print(f"Текущие аргументы: {args}")
                        print("Введите изменения в формате key=value через запятую.")
                        raw = input("Новые аргументы: ").strip()
                        new_args = dict(args)
                        for part in raw.split(","):
                            part = part.strip()
                            if "=" in part:
                                k, v = part.split("=", 1)
                                new_args[k.strip()] = v.strip()
                        decisions.append({
                            "type": "edit",
                            "edited_action": {"name": name, "args": new_args},
                        })
                        break

                    else:
                        print(
                            f"Недопустимый выбор. Варианты: {', '.join(hints)}"
                        )

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

        # ── Финальный ответ ───────────────────────────────────────────────
        # Сообщения — объекты LangChain (AIMessage и др.),
        # текст хранится в атрибуте .content, а не в ключе словаря.
        messages = result.get("messages", [])
        if messages:
            print(f"\nАгент: {messages[-1].content}")


if __name__ == "__main__":
    main()