import sys
from agent import graph

def main():
    print("LangGraph stream‑mode demo. Type 'exit' to quit.")
    state = {"messages": []}
    # Если передан аргумент, используем его как сообщение; иначе читаем из stdin
    if len(sys.argv) > 1:
        user_input = " ".join(sys.argv[1:])
        process(user_input, state)
    else:
        while True:
            try:
                user_input = input("You: ")
            except EOFError:
                break
            if not user_input:
                continue
            if user_input.lower() in {"exit", "quit"}:
                break
            process(user_input, state)

def process(user_input, state):
    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]
            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"]

if __name__ == "__main__":
    main()