# Human‑in‑the‑Loop (HITL) middleware example
# interrupt decision types: approve, edit, reject, respond
# respond decision can be used to supply a human reply as tool result

import os
from typing import List, Dict

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

# --- Tool ---------------------------------------------------------------

class GetWeatherTool(BaseTool):
    name: str = "get_weather"
    description: str = "Get the weather for a city on a given date."
    def _run(self, city: str, date: str) -> str:
        return f"The weather in {city} on {date} is sunny with a high of 25°C."

# --- LLM ---------------------------------------------------------------
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    raise RuntimeError("OPENAI_API_KEY environment variable is required.")

llm = ChatOpenAI(api_key=api_key, temperature=0.7)

# --- Agent --------------------------------------------------------------
memory = MemorySaver()

agent = create_agent(
    model=llm,
    tools=[GetWeatherTool()],
    system_prompt="You are a helpful assistant.",
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={
                "get_weather": True,  # allow approve, edit, reject
            },
            description_prefix="Подтвердите вызов инструмента",
        ),
    ],
    checkpointer=memory,
)

# --- Interactive loop -----------------------------------------------
if __name__ == "__main__":
    config = {"configurable": {"thread_id": "session-1"}}
    while True:
        try:
            try:
                user_input = input("\nUser: ")
            except EOFError:
                print("\nNo input provided. Exiting.")
                break
            if user_input.lower() in {"exit", "quit", "q"}:
                print("Goodbye!")
                break

            # First invocation
            result = agent.invoke({"messages": [{"role": "human", "content": user_input}]}, config=config)

            # Process interrupts until finished
            while "__interrupt__" in result:
                interrupt = result["__interrupt__"][0].value
                action_requests = interrupt.get("action_requests", [])
                review_configs = interrupt.get("review_configs", [])

                decisions: List[Dict] = []
                print("\n--- Подтверждение ---")
                for idx, action in enumerate(action_requests):
                    name = action.get("name")
                    args = action.get("args", {})
                    description = action.get("description", "")
                    print(f"Инструмент: {name}")
                    print(f"Аргументы: {args}")
                    if description:
                        print(f"Описание: {description}")
                    # Determine allowed decisions
                    allowed = ["approve", "reject"]
                    # Check if edit is allowed
                    for cfg in review_configs:
                        if cfg.get("name") == name:
                            allowed = cfg.get("allowed_decisions", allowed)
                            break
                    # Prompt user
                    while True:
                        choice = input("a = approve, r = reject: ").strip().lower()
                        if choice == "a" and "approve" in allowed:
                            decisions.append({"type": "approve"})
                            break
                        if choice == "r" and "reject" in allowed:
                            msg = input("Сообщение для агента (причина отказа): ")
                            decisions.append({"type": "reject", "message": msg})
                            break
                        print("Неверный ввод. Попробуйте снова.")

                # Resume
                result = agent.invoke(Command(resume={"decisions": decisions}), config=config)  # Command(resume={"decisions":[…]})

            # Finished
            final_message = result.get("messages", [])[-1].get("content", "")
            print(f"\nAssistant: {final_message}\n")

        except KeyboardInterrupt:
            print("\nInterrupted. Exiting.")
            break