import os
from typing import Any, Dict, List

# ────────────────────────  Imports from the required stack  ────────────────────────
from langchain_ollama import ChatOllama
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import MemorySaver

# ────────────────────────  LLM and tool definition  ─────────────────────────────
llm = ChatOllama(model="llama3")

def get_weather(city: str, date: str) -> str:
    """Mock implementation of a weather‑lookup tool."""
    return f"Погода в {city} на {date}: солнечно, +25°C"

# ────────────────────────  Agent creation with HumanInTheLoopMiddleware  ───────
memory = MemorySaver()

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

# ────────────────────────  Helper for the human‑in‑the‑loop loop  ───────────────
def run_agent_with_human_loop(user_message: str, thread_id: str = "session-1") -> None:
    """
    Запускает агента с возможностью подтверждения вызова инструментов.
    user_message – сообщение пользователя (строка).
    thread_id   – идентификатор сессии для сохранения истории.
    """
    from langgraph.types import Command

    config = {"configurable": {"thread_id": thread_id}}

    # первый запуск
    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["action_requests"]
        review_configs = interrupt_value["review_configs"]

        decisions: List[Dict[str, Any]] = []

        print("\n--- Подтверждение ---")
        for idx, (req, cfg) in enumerate(zip(action_requests, review_configs), start=1):
            name = req.get("name")
            args = req.get("args", {})
            description = req.get("description")

            print(f"\n{idx}. Инструмент: {name}")
            print(f"   Аргументы: {args}")
            if description:
                print(f"   Описание: {description}")

            # определяем доступные решения
            allowed = cfg.get("allowed_decisions", ["approve", "edit", "reject"])
            options = ", ".join(
                [f"{'a'=approve}" if d == "approve" else f"{'r'=reject}"
                 for d in allowed]
            )
            # простейший ввод (без проверки)
            choice = input(f"a = approve, r = reject: ").strip().lower()
            if choice == "a":
                decisions.append({"type": "approve"})
            elif choice == "r":
                msg = input("Сообщение для агента (причина отказа): ")
                decisions.append({"type": "reject", "message": msg})
            else:
                # если пользователь ввёл что‑то другое, считаем reject
                decisions.append({"type": "reject", "message": "неизвестное решение"})

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

    # вывод финального ответа
    final_message = result["messages"][-1]["content"]
    print("\nАгент:", final_message)


# ────────────────────────  Пример использования  ───────────────────────────────
if __name__ == "__main__":
    user_input = input("Вы: ")
    run_agent_with_human_loop(user_input)