"""
Задание 5: Human-in-the-Loop через middleware

Агент с HumanInTheLoopMiddleware: при каждом вызове инструмента
агент останавливается, пользователь подтверждает (approve/reject),
после чего выполнение возобновляется.
"""

import os
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


# ─── Инструменты ─────────────────────────────────────────────────────────────

@tool
def get_weather(city: str, date: str = "сегодня") -> str:
    """
    Получить информацию о погоде в городе.
    Вход: city — название города, date — дата (по умолчанию 'сегодня').
    """
    # Заглушка — реальное обращение к API погоды
    return f"В городе {city} на {date} ожидается переменная облачность, +15°C."


@tool
def search_web(query: str) -> str:
    """
    Поиск информации в интернете.
    Вход: query — поисковый запрос.
    """
    return f"Результаты поиска по запросу '{query}': [результат-1, результат-2]"


TOOLS = [get_weather, search_web]


# ─── Инициализация LLM ───────────────────────────────────────────────────────

api_base = os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1")
api_key = os.environ.get("OPENAI_API_KEY", "")
model_name = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")

llm = ChatOpenAI(
    model=model_name,
    openai_api_base=api_base,
    openai_api_key=api_key,
    temperature=0,
)


# ─── Сборка агента с HumanInTheLoopMiddleware ────────────────────────────────

memory = MemorySaver()

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


# ─── Цикл подтверждения ─────────────────────────────────────────────────────

def get_user_decisions(action_requests: list[dict], review_configs: list[dict]) -> list[dict]:
    """
    Показать пользователю запросы на вызов инструментов и собрать решения.
    Поддерживает approve и reject.
    """
    decisions = []

    for i, (action, review_cfg) in enumerate(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{'─' * 50}")
        print(f"  Подтверждение #{i + 1}")
        print(f"{'─' * 50}")
        print(f"  Инструмент : {name}")
        print(f"  Аргументы   : {json.dumps(args, ensure_ascii=False)}")
        if description:
            print(f"  Описание   : {description}")
        print(f"  Решения     : {', '.join(allowed)}")

        while True:
            choice = input(f"\n  (a)pprove / (r)eject: ").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:
                msg = input("  Сообщение для агента (причина отказа): ").strip()
                decisions.append({
                    "type": "reject",
                    "message": msg or "Запрос отклонён пользователем",
                })
                break
            else:
                print(f"  Неверный ввод. Доступные: {', '.join(allowed)}")

    return decisions


def run_agent_with_hitl(user_message: str, thread_id: str = "session-1"):
    """
    Запустить агента с обработкой Human-in-the-Loop.
    Цикл: invoke → проверка __interrupt__ → запрос решений → resume → ...
    """
    config = {"configurable": {"thread_id": thread_id}}

    print("=" * 60)
    print("  🤖 Агент с Human-in-the-Loop Middleware")
    print("=" * 60)
    print(f"\n  Вы: {user_message}\n")

    result = agent.invoke(
        {"messages": [{"role": "human", "content": user_message}]},
        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,
        )

    # Финальный ответ
    final_message = result["messages"][-1]
    print(f"\n{'=' * 60}")
    print(f"  Агент: {final_message.content}")
    print(f"{'=' * 60}")

    return result


# ─── Точка входа ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys
    message = sys.argv[1] if len(sys.argv) > 1 else "Какая погода в Казани сегодня?"
    run_agent_with_hitl(message)