"""
Human-in-the-Loop через middleware
Использует: langchain >= 0.3, langgraph >= 0.2
Установка: pip install langchain langchain-openai langgraph
"""

import os
import sys
from typing import Any

# ── Импорты ──────────────────────────────────────────────────────────────────
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 InMemorySaver   # InMemorySaver, не MemorySaver
from langgraph.types import Command

# ── Модель ───────────────────────────────────────────────────────────────────
# Ключ читается из переменной окружения OPENAI_API_KEY
if not os.environ.get("OPENAI_API_KEY"):
    print(
        "Ошибка: переменная окружения OPENAI_API_KEY не задана.\n"
        "Задайте её командой: export OPENAI_API_KEY=ваш-ключ",
        file=sys.stderr,
    )
    sys.exit(1)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

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

# ── Агент с HumanInTheLoopMiddleware ─────────────────────────────────────────
memory = InMemorySaver()   # обязателен для сохранения паузы

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

# ── Обработка одного прерывания ───────────────────────────────────────────────
def handle_interrupt(interrupts) -> list[dict[str, Any]]:
    """
    Показывает пользователю запросы на подтверждение и собирает список решений.

    interrupts — кортеж объектов Interrupt из result.interrupts (GraphOutput, version="v2").
    Возвращает список решений в том же порядке, что и action_requests.
    """
    interrupt_value: dict = interrupts[0].value
    action_requests: list = interrupt_value.get("action_requests", [])
    review_configs: list  = interrupt_value.get("review_configs", [])

    decisions: list[dict[str, Any]] = []

    for idx, action in enumerate(action_requests):
        name: str        = action.get("name", "")
        # В v2-документации поле называется «arguments», а не «args»
        arguments: dict  = action.get("arguments", action.get("args", {}))
        description: str = action.get("description", "")

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

        # Определяем допустимые решения для этого действия
        allowed: list[str] = ["approve", "reject"]
        for cfg in review_configs:
            # Поле называется «action_name», а не «name»
            if cfg.get("action_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"Текущие аргументы: {arguments}")
                print("Введите новые аргументы в формате ключ=значение (например: city=Москва).")
                new_args = dict(arguments)  # копируем оригинал
                raw = input("Новые аргументы: ").strip()
                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)}")

    return decisions


# ── Основная логика: запрос → пауза → цикл resume ────────────────────────────
def ask_and_run(user_text: str, config: dict) -> None:
    """
    Отправляет сообщение агенту, обрабатывает все прерывания HITL
    и выводит финальный ответ.
    """
    # Первый вызов агента
    # version="v2" — обязателен для получения GraphOutput с .interrupts
    result = agent.invoke(
        {"messages": [{"role": "human", "content": user_text}]},
        config=config,
        version="v2",
    )

    # Цикл: пока агент хочет подтверждения — спрашиваем и возобновляем
    while result.interrupts:
        decisions = handle_interrupt(result.interrupts)

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

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


# ── Чат-цикл ─────────────────────────────────────────────────────────────────
def main() -> None:
    # 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

        ask_and_run(user_text, config)


if __name__ == "__main__":
    main()