import sys
from agent import graph

if __name__ == "__main__":
    print("LangGraph stream‑mode demo. Type 'exit' to quit.")
    state = {"messages": []}
    for line in sys.stdin:
        user_input = line.rstrip('\n')
        if not user_input:
            continue
        if user_input.lower() in {"exit", "quit"}:
            break
        # Stream output with separators between steps
        current_step = None
        for chunk_type, chunk_data in graph.stream({"input": user_input, "messages": state["messages"]}, stream_mode=["messages", "updates"]):
            if chunk_type == "messages":
                message, meta = chunk_data
                step = meta.get("langgraph_step")
                if current_step is None:
                    current_step = step
                elif step != current_step:
                    print("\n--- --- ---\n", end="")
                    current_step = step
                if message.content:
                    print(message.content, end="")
            elif chunk_type == "updates":
                last_message = chunk_data.get("model", {}).get("messages", [-1])[-1]
                # If the last message is an AIMessage with content, print it
                try:
                    if hasattr(last_message, "content") and last_message.content:
                        print(last_message.content, end="")
                except Exception:
                    pass
        print()
        state["messages"] = graph.invoke({"input": user_input, "messages": state["messages"]})["messages"]
    print("Goodbye!")