import json
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from langchain.tools import tool

# Define a simple get_weather tool
@tool
def get_weather(city: str, date: str = "today") -> str:
    """Return a mock weather for the given city and date."""
    return f"The weather in {city} on {date} is sunny with a high of 25°C."

# Initialize LLM (replace with your OpenRouter key if needed)
llm = ChatOpenAI(model="gpt-4o", temperature=0)

memory = MemorySaver()

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

if __name__ == "__main__":
    import sys
    from langgraph.types import Command

    # Simple CLI loop
    thread_id = "session-1"
    config = {"configurable": {"thread_id": thread_id}}

    # Initial user message
    user_msg = " ".join(sys.argv[1:]) or "Какая погода в Казани сегодня?"
    result = agent.invoke({"messages": [{"role": "human", "content": user_msg}]}, config=config)

    while "__interrupt__" in result:
        interrupt = result["__interrupt__"][0].value
        action_requests = interrupt["action_requests"]
        review_configs = interrupt["review_configs"]
        decisions = []
        for idx, action in enumerate(action_requests):
            print(f"\n--- Подтверждение ---")
            print(f"Инструмент: {action['name']}")
            print(f"Аргументы: {action['args']}")
            if "description" in action:
                print(f"Описание: {action['description']}")
            # Determine allowed decisions
            allowed = review_configs.get(action['name'], {}).get("allowed_decisions", ["approve", "reject", "edit"]) if isinstance(review_configs.get(action['name']), dict) else ["approve", "reject", "edit"]
            # Simple input handling
            while True:
                inp = input("a = approve, r = reject, e = edit: ").strip().lower()
                if inp == "a" and "approve" in allowed:
                    decisions.append({"type": "approve"})
                    break
                if inp == "r" and "reject" in allowed:
                    msg = input("Введите причину отказа: ")
                    decisions.append({"type": "reject", "message": msg})
                    break
                if inp == "e" and "edit" in allowed:
                    # For simplicity, ask for new args as JSON
                    new_args = input("Введите новые аргументы в формате JSON: ")
                    try:
                        new_args_dict = json.loads(new_args)
                        decisions.append({"type": "edit", "edited_action": {"name": action['name'], "args": new_args_dict}})
                        break
                    except Exception as exc:
                        print("Неверный JSON, попробуйте снова.")
                print("Недопустимый ввод. Попробуйте снова.")
        # Resume
        result = agent.invoke(Command(resume={"decisions": decisions}), config=config)

    # Final answer
    final_msg = result["messages"][-1]["content"]
    print("\nАгент: " + final_msg)