#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
solution.py

Пример использования HumanInTheLoopMiddleware в LangChain Agent.
При каждом вызове инструмента агент останавливается и запрашивает у пользователя подтверждение
(approve / reject). После принятия решения выполнение возобновляется через Command.

Требования:
    pip install langchain langgraph openai rich
"""

import json
from typing import Any, Dict, List

# --------------------------------------------------------------------------- #
# 1. Подключаем необходимые модули LangChain и LangGraph
# --------------------------------------------------------------------------- #

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

# --------------------------------------------------------------------------- #
# 2. Определяем простой инструмент get_weather
# --------------------------------------------------------------------------- #

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

# --------------------------------------------------------------------------- #
# 3. Создаём LLM (используем OpenAI GPT-4o-mini как пример)
# --------------------------------------------------------------------------- #

from langchain_openai import ChatOpenAI

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

# --------------------------------------------------------------------------- #
# 4. Настраиваем агент с HumanInTheLoopMiddleware
# --------------------------------------------------------------------------- #

memory = MemorySaver()

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

# --------------------------------------------------------------------------- #
# 5. Функция для получения решений от пользователя
# --------------------------------------------------------------------------- #

def ask_decisions(action_requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Для каждого запроса к инструменту выводим информацию и запрашиваем у пользователя решение.
    Возвращаем список словарей с решениями в том же порядке, что и action_requests.
    """
    decisions = []
    print("\n--- Подтверждение вызова инструмента ---")
    for idx, req in enumerate(action_requests):
        name = req.get("name", "неизвестный инструмент")
        args = req.get("args", {})
        description = req.get("description", "")

        print(f"\n{idx + 1}. Инструмент: {name}")
        if description:
            print(f"   Описание: {description}")
        print(f"   Аргументы: {json.dumps(args, ensure_ascii=False)}")

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

# --------------------------------------------------------------------------- #
# 6. Основной цикл взаимодействия с агентом
# --------------------------------------------------------------------------- #

def main() -> None:
    """
    Запускает чат-цикл: пользователь вводит сообщение,
    агент обрабатывает его, при необходимости запрашивает подтверждение.
    После завершения выводится финальный ответ агента.
    """
    thread_id = "session-1"
    config = {"configurable": {"thread_id": thread_id}}

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

    while True:
        user_msg = input("\nВы: ").strip()
        if user_msg.lower() in ("выход", "quit", "exit"):
            print("До свидания!")
            break

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

        # Цикл подтверждений
        while "__interrupt__" in result:
            interrupt_value = result["__interrupt__"][0].value
            action_requests = interrupt_value.get("action_requests", [])
            # review_configs не используется в этом примере, но можно вывести при желании

            decisions = ask_decisions(action_requests)

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

        # После завершения выводим ответ агента
        if "messages" in result and result["messages"]:
            last_msg = result["messages"][-1]
            print(f"\nАгент: {last_msg.get('content', '')}")
        else:
            print("\nАгент не вернул ответа.")

if __name__ == "__main__":
    main()