"""
Human‑in‑the‑loop demo using LangChain + LangGraph middleware.

Run with:
    python main.py

The script creates an agent that asks the user to approve or reject each tool call.
The user interacts via the terminal.
"""

import os
import json
from typing import List, Dict, Any

# LangChain imports
from langchain_openai import ChatOpenAI
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

# Simple tool: get_weather

def get_weather(city: str, date: str) -> str:
    """Return a dummy weather report.

    In a real application this would call an external API.
    """
    return f"Погода в {city} на {date}: солнечно, 25°C."

# Build the agent

def build_agent() -> Any:
    # LLM – replace with your own key / model if needed
    llm = ChatOpenAI(temperature=0.0)

    memory = MemorySaver()

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

# Helper to pretty‑print action requests

def show_action_requests(action_requests: List[Dict[str, Any]]) -> None:
    print("\n--- Подтверждение ---")
    for idx, act in enumerate(action_requests, start=1):
        name = act.get("name")
        args = act.get("args")
        description = act.get("description")
        print(f"{idx}. Инструмент: {name}")
        print(f"   Аргументы: {json.dumps(args, ensure_ascii=False)}")
        if description:
            print(f"   Описание: {description}")
    print()

# Ask user for decisions

def ask_decisions(action_requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    decisions: List[Dict[str, Any]] = []
    for act in action_requests:
        while True:
            inp = input("a = approve, r = reject: ").strip().lower()
            if inp == "a":
                decisions.append({"type": "approve"})
                break
            elif inp == "r":
                msg = input("Сообщение для агента (причина отказа): ").strip()
                decisions.append({"type": "reject", "message": msg})
                break
            else:
                print("Неверный ввод. Попробуйте снова.")
    return decisions

# Main loop

def run_agent(agent: Any) -> None:
    thread_id = "session-1"
    config = {"configurable": {"thread_id": thread_id}}

    # Initial user message
    user_msg = input("Вы: ")
    messages = [{"role": "human", "content": user_msg}]

    # First invoke
    result = agent.invoke({"messages": messages}, config=config)

    # Loop until no interrupt
    while "__interrupt__" in result:
        interrupt = result["__interrupt__"][0].value
        action_requests = interrupt.get("action_requests", [])
        # review_configs = interrupt.get("review_configs", [])  # not used here
        show_action_requests(action_requests)
        decisions = ask_decisions(action_requests)
        # Resume
        result = agent.invoke(Command(resume={"decisions": decisions}), config=config)

    # Final answer
    final_msg = result.get("messages", [])[-1].get("content", "")
    print(f"\nАгент: {final_msg}")

if __name__ == "__main__":
    # Ensure OpenAI key is set if using ChatOpenAI
    if os.getenv("OPENAI_API_KEY") is None:
        print("WARNING: OPENAI_API_KEY not set. Using default model may fail.")
    agent = build_agent()
    run_agent(agent)