"""Console entry point that prints LangGraph agent output in stream mode."""

from __future__ import annotations

import sys
from typing import Any

from agent import agent


_current_step: int | None = None


def format_message(message: Any) -> str:
    """Format a completed model message or a tool-call request."""

    content = getattr(message, "content", "")
    if content:
        return str(content)

    tool_calls = getattr(message, "tool_calls", None) or []
    if tool_calls:
        call = tool_calls[0]
        return f"{call['name']}({call['args']})"

    return ""


def format_chunk_message(chunk: tuple[Any, dict[str, Any]]) -> None:
    """Print token chunks and separate LangGraph steps."""

    global _current_step

    message, meta = chunk
    langgraph_step = meta.get("langgraph_step")
    if _current_step is None:
        _current_step = langgraph_step
    elif langgraph_step != _current_step:
        _current_step = langgraph_step
        print("\n--- --- ---\n", end="")

    if message.content:
        print(message.content, end="", flush=True)


def run_stream(query: str) -> None:
    """Run the agent with .stream() and handle messages/updates chunks."""

    global _current_step
    _current_step = None

    stream = agent.stream(
        {"messages": [{"role": "human", "content": query}]},
        stream_mode=["messages", "updates"],
    )

    for chunk in stream:
        chunk_type, chunk_data = chunk

        if chunk_type == "messages":
            format_chunk_message(chunk_data)

        if chunk_type == "updates":
            model_update = chunk_data.get("model")
            if model_update:
                last_message = model_update["messages"][-1]
                formatted = format_message(last_message)
                if formatted:
                    print(formatted, end="", flush=True)

    print()


def main() -> None:
    """Read one CLI query or start a small interactive loop."""

    if len(sys.argv) > 1:
        run_stream(" ".join(sys.argv[1:]))
        return

    print("Stream-mode AI agent. Type 'exit' to quit.")
    while True:
        try:
            query = input("You: ").strip()
        except EOFError:
            print()
            break

        if query.lower() in {"exit", "quit"}:
            break
        if query:
            run_stream(query)


if __name__ == "__main__":
    main()