from langchain_core.messages import HumanMessage
from .graph import graph

# -> return updates from each Node as it finishes
# for update in graph.stream(
#     initial_state,
#     stream_mode="updates"
# ):
#     print(update)

# stream("updates") → Each Node's State updates.
# stream("values") → Full accumulated State at each step.

def main():
    config = {
        "configurable": {
            "thread_id": "restaurant-test-1"
        }
    }

    # First user message
    result = graph.invoke(
        {
            "messages": [
                HumanMessage(
                    content="Book a table at Olive Garden for 4 people"
                )
            ],
            "global_params": {},
            "selected_tool": None,
            "missing_params": [],
            "tool_args": {},
            "tool_result": None
        },
        config=config
    )

    print("ASSISTANT:", result["messages"][-1].content)

    # Second user message
    result = graph.invoke(
        {
            "messages": [
                HumanMessage(content="Tomorrow at 7 PM")
            ]
        },
        config=config
    )

    print("SELECTED TOOL:", result["selected_tool"])
    print("TOOL ARGS:", result["tool_args"])
    print("TOOL RESULT:", result["tool_result"])
    print("ASSISTANT:", result["messages"][-1].content)

    # Generate Graph image
    with open("agent_graph.png", "wb") as f:
        f.write(graph.get_graph().draw_mermaid_png())

    print("Graph saved to agent_graph.png")


if __name__ == "__main__":
    main()