import asyncio
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph
from langgraph.types import Command

# LLM
llm = ChatOpenAI(model="gpt-4o-mini", base_url="https://openrouter.ai/api/v1", api_key=os.getenv("OPENAI_API_KEY"), temperature=0.0)

# Simple tool
@tool
def get_weather(city: str, date: str) -> str:
    return f"Погода в {city} на {date}: солнечно, 25°C"

# Agent with middleware
memory = MemorySaver()

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

async def main():
    config = {"configurable": {"thread_id": "сессия-1"}}
    result = await agent.ainvoke(
        {"messages": [{"role": "human", "content": "Какая погода в Казани сегодня?"}]},
        config=config,
    )
    # loop for interrupts
    while "__interrupt__" in result:
        interrupt = result["__interrupt__"][0].value
        decisions = []
        for req in interrupt["action_requests"]:
            print(f"Инструмент: {req['name']}")
            print(f"Аргументы: {req['args']}")
            if "description" in req:
                print(req["description"])
            choice = input("a=approve, r=reject: ")
            if choice.lower() == "a":
                decisions.append({"type": "approve"})
            else:
                msg = input("Причина отказа: ")
                decisions.append({"type": "reject", "message": msg})
        result = await agent.ainvoke(Command(resume={"decisions": decisions}), config=config)
    print("\nОтвет: ", result["messages"][-1]["content"])

if __name__ == "__main__":
    asyncio.run(main())