from llama_index.core.agent.workflow import (
    AgentInput,
    AgentOutput,
    ToolCall,
    ToolCallResult,
    AgentStream,
)

async def run_workflow_agent(workflow_agent, user_msg: str):
    handler = workflow_agent.run(user_msg=user_msg)

    current_agent = None

    async for event in handler.stream_events():
        # Detect agent switch
        if hasattr(event, "current_agent_name") and event.current_agent_name != current_agent:
            current_agent = event.current_agent_name
            print(f"\n{'='*50}")
            print(f"Agent Name: {current_agent}")
            print(f"{'='*50}\n")

        # Handle Agent Output
        if isinstance(event, AgentOutput):
            if event.response.content:
                print("Output:", event.response.content)
            if event.tool_calls:
                print("\nPlanning to use tools:", [call.tool_name for call in event.tool_calls])

        # Handle Tool Call
        elif isinstance(event, ToolCall):
            print("\n**************************Calling***********************")
            print(f"Calling Tool: {event.tool_name}")
            print(f"Calling With Arguments: {event.tool_kwargs}")

        # Handle Tool Call Result
        elif isinstance(event, ToolCallResult):
            print("\n**************************Results***********************")
            print(f"Tool Result ({event.tool_name}):")
            print(f"Arguments: {event.tool_kwargs}")
            print(f"Output: {event.tool_output}")
            print(f"{'='*100}")

    return handler