#!/usr/bin/env python3
"""Stream an agent chat turn: print reply text and tool calls as they happen.

Defaults to wfref-demo-responder (fast smoke). Pass --agent <name> for a
specialist, or --workflow <name> to run a workflow instead. System agents such
as aiden-router accept chat even when they are absent from --list.

Usage:
  export STACKGEN_URL=https://app.stackgen.com
  export STACKGEN_TOKEN=stackgen_…
  export STACKGEN_WORKSPACE_ID=<workspace-uuid>

  # Discover agents and workflows in the workspace
  python examples/ask/stream-agent-chat.py --list

  # Stream a short reply (agent chat)
  python examples/ask/stream-agent-chat.py \
    --agent wfref-demo-responder \
    --message "Reply with exactly: pong"

  # Or start a workflow run and stream the same AG-UI events
  python examples/ask/stream-agent-chat.py \
    --workflow incident-triage \
    --message "Summarize checkout latency"

The generator stops on RUN_FINISHED or when max_wall_clock_ms elapses.
"""

from __future__ import annotations

import argparse
import os
import sys

from stackgen import StackgenClient, StackgenConfig


def require_env(name: str, *fallbacks: str) -> str:
    for key in (name, *fallbacks):
        value = (os.environ.get(key) or "").strip()
        if value:
            return value
    raise SystemExit(f"missing env {name}")


def text_field(data: dict, *keys: str) -> str:
    for key in keys:
        value = data.get(key)
        if isinstance(value, str) and value:
            return value
    return ""


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Stream AG-UI events for an agent chat or workflow run.",
    )
    parser.add_argument("--message", default="Reply with exactly: pong")
    parser.add_argument(
        "--agent",
        default="wfref-demo-responder",
        help="Agent name for start_agent_chat (ignored when --workflow is set)",
    )
    parser.add_argument(
        "--workflow",
        default="",
        help="Workflow name for run_workflow instead of agent chat",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="Print list_agents and list_workflows, then exit",
    )
    parser.add_argument("--timeout-seconds", type=int, default=300)
    args = parser.parse_args()

    client = StackgenClient(
        StackgenConfig(
            base_url=require_env("STACKGEN_URL"),
            api_token=require_env("STACKGEN_TOKEN"),
            workspace_id=require_env("STACKGEN_WORKSPACE_ID", "STACKGEN_PROJECT"),
        )
    )

    if args.list:
        agents = client.aiden.list_agents()
        workflows = client.aiden.list_workflows(latestOnly=True, limit=50)
        print("agents:")
        for a in agents[:30]:
            print(f"  - {text_field(a, 'name') or a}")
        print("workflows:")
        for w in workflows[:30]:
            print(f"  - {text_field(w, 'name') or w}")
        return

    if args.workflow:
        started = client.aiden.run_workflow(
            args.workflow, {"input": {"message": args.message}}
        )
        label = f"workflow={args.workflow}"
    else:
        started = client.aiden.start_agent_chat(
            args.agent, {"message": args.message}
        )
        label = f"agent={args.agent}"

    session_id = text_field(started, "session_id", "sessionId")
    trace_id = text_field(started, "trace_id", "traceId")
    print(f"{label} session_id={session_id} trace_id={trace_id}")
    if not trace_id:
        raise SystemExit("start did not return trace_id")

    for ev in client.aiden.stream_execution_events(
        trace_id,
        max_wall_clock_ms=args.timeout_seconds * 1000,
    ):
        etype = ev.get("type")
        data = ev.get("data") if isinstance(ev.get("data"), dict) else {}
        if etype == "TEXT_MESSAGE_CONTENT":
            chunk = text_field(data, "Content", "content", "delta", "Delta")
            if chunk:
                sys.stdout.write(chunk)
                sys.stdout.flush()
        elif etype == "TOOL_CALL_START":
            name = text_field(data, "ToolCallName", "toolCallName", "name")
            tid = text_field(data, "ToolCallID", "toolCallId", "id")
            print(f"\n[tool start] {name or 'tool'} ({tid or '?'})")
        elif etype == "TOOL_CALL_ARGS":
            delta = text_field(data, "delta", "Delta", "Content", "content")
            if delta:
                sys.stdout.write(delta)
                sys.stdout.flush()
        elif etype in ("TOOL_CALL_RESULT", "TOOL_CALL_END"):
            print(f"\n[tool {'result' if etype == 'TOOL_CALL_RESULT' else 'end'}]")
        elif etype == "RUN_ERROR":
            print(f"\n[run error] {data}", file=sys.stderr)
        elif etype == "RUN_FINISHED":
            print("\n[run finished]")


if __name__ == "__main__":
    main()