# imports from langgraph
from langgraph.graph import START, END

# imports from local files
from src.graph_builder import graph_builder
from src.nodes import chatbot, BasicToolNode
from src.utils import route_tools
from src.tools import tool


def main():
    # The first argument is the unique node name
    # The second argument is the function or object that will be called whenever
    # the node is used.
    graph_builder.add_node("chatbot", chatbot)      # Create the chatbot node
    graph_builder.add_edge("chatbot", END)          # Add the edge from chatbot node to END
    tool_node = BasicToolNode(tools=[tool])         # Intermediate variable creation to create a tool node with the tools
    graph_builder.add_node("tools", tool_node)      # Create a tool node with the tools
    graph_builder.add_conditional_edges(            # Create the conditionnal edges from chatbot to tools
        "chatbot",
        route_tools,
        {"tools": "tools", END: END},
        )
    graph_builder.add_edge("tools", "chatbot")      # Add the edge from tools node back to chatbot
    graph_builder.add_edge(START, "chatbot")        # Add the edge from START to chatbot node
    graph = graph_builder.compile()
    
    def stream_graph_updates(user_input: str):
        for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
            for value in event.values():
                print("Assistant:", value["messages"][-1].content)


    while True:
        try:
            user_input = input("User: ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("Goodbye!")
                break
            stream_graph_updates(user_input)
        except:
            # fallback if input() is not available
            user_input = "What do you know about LangGraph?"
            print("User: " + user_input)
            stream_graph_updates(user_input)
            break

if __name__ == "__main__":
    main()