"""Console demo for stream-mode LangGraph agent."""

import os
from typing import Any

from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph

from agent import agent

# Helper to format messages

def format_message(message: Any) -> str:
    if hasattr(message, "content") and message.content:
        return message.content
    # fallback for tool calls
    if hasattr(message, "tool_calls") and message.tool_calls:
        call = message.tool_calls[0]
        return f"{call['name']}({call['args']})"
    return ""

# Stream the conversation

def main() -> None:
    # initial human message
    human_msg = HumanMessage(content="Составь список покупок: молоко, хлеб, яблоки")
    # start stream
    stream = agent.stream(
        {"messages": [human_msg]},
        stream_mode=["messages", "updates"],
    )

    step = 1
    for chunk in stream:
        chunk_type, chunk_data = chunk
        if chunk_type == "messages":
            message, meta = chunk_data
            if meta.get("langgraph_step") != step:
                step = meta.get("langgraph_step")
                print("\n--- --- ---\n")
            if message.content:
                print(message.content, end="", flush=True)
        elif chunk_type == "updates":
            if chunk_data.get("model"):
                last_msg = chunk_data["model"]["messages"][-1]
                print(format_message(last_msg))

if __name__ == "__main__":
    main()