import os
from uuid import UUID

from dotenv import load_dotenv

from swarmd_sdk import SwarmDClient

load_dotenv()

me = UUID(os.environ["SWARMD_AGENT_ID"])

with SwarmDClient.from_env() as client:
    for agent in client.get_agent_subscriptions(me):
        print(f"{agent.name}")
        print(f"  {agent.description}")
        print(f"  card: {agent.agent_card_url}")

from dotenv import load_dotenv

from swarmd_sdk import McpClient

load_dotenv()

with McpClient.from_env() as mcp:
    servers = mcp.list_available()
    if not servers:
        raise SystemExit("no MCP servers granted to this agent yet")

    server = servers[0]
    print(f"{server.name} ({server.mcp_server_id})")

    for tool in mcp.list_tools(server.mcp_server_id):
        print(f"  - {tool['name']}: {tool.get('description', '')}")

    result = mcp.call_tool(
        server.mcp_server_id,
        "search",
        {"query": "weather in Tokyo"},
    )
    print(result)

import os
from uuid import UUID

import uvicorn
from dotenv import load_dotenv
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route

from swarmd_sdk import CorrelationIdMiddleware, SwarmDRuntime, create_admin_app

load_dotenv()

# 1. Build and configure the runtime.
runtime = SwarmDRuntime()
runtime.configure(
    agent_id=os.environ["SWARMD_AGENT_ID"],
    client_secret=os.environ["SWARMD_CLIENT_SECRET"],
    base_url=os.getenv("SWARMD_BASE_URL", "https://api.dev.swarmd.ai"),
    token_url=os.environ["SWARMD_TOKEN_URL"],
    webhook_secret=os.getenv("SWARMD_WEBHOOK_SECRET"),
)

# 2. Your catalogue, discovered at boot and rebuilt on refresh.
CATALOGUE = {"agents": (), "mcp_servers": ()}


def refresh() -> None:
    """Called at boot, on POST /admin/refresh, and on every signed webhook."""
    me = UUID(runtime.config.agent_id)
    CATALOGUE["agents"] = tuple(runtime.client.get_agent_subscriptions(me))
    CATALOGUE["mcp_servers"] = tuple(runtime.mcp.list_available())
    print(
        f"[refresh] {len(CATALOGUE['agents'])} agent(s), "
        f"{len(CATALOGUE['mcp_servers'])} MCP server(s)"
    )


refresh()


# 3. Your agent's own endpoints. Replace with a real A2A server in production.
async def agent_card(request):
    return JSONResponse({
        "name": "my_agent",
        "description": "An agent built on the core Swarmd SDK",
        "url": f"http://localhost:{os.getenv('PORT', '8080')}/",
        "version": "0.1.0",
        "capabilities": {"streaming": False},
        "defaultInputModes": ["text/plain"],
        "defaultOutputModes": ["text/plain"],
        "skills": [],
    })


async def health(request):
    return JSONResponse({"status": "ok"})


app = Starlette(routes=[
    Route("/.well-known/agent-card.json", agent_card),
    Route("/health", health),
])

# 4. Correlation propagation, then the admin surface.
app.add_middleware(CorrelationIdMiddleware)
app.mount("/admin", create_admin_app(runtime, on_refresh=refresh))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "8080")))

import asyncio
import os
from uuid import UUID, uuid4

import httpx
from dotenv import load_dotenv

from swarmd_sdk import SwarmDRuntime

load_dotenv()

runtime = SwarmDRuntime()
runtime.configure(
    agent_id=os.environ["SWARMD_AGENT_ID"],
    client_secret=os.environ["SWARMD_CLIENT_SECRET"],
    base_url=os.getenv("SWARMD_BASE_URL", "https://api.dev.swarmd.ai"),
    token_url=os.environ["SWARMD_TOKEN_URL"],
)

TERMINAL = {"completed", "failed", "canceled", "rejected", "unknown"}


async def ask(agent_name: str, prompt: str) -> dict:
    me = UUID(runtime.config.agent_id)
    target = next(
        a for a in runtime.client.get_agent_subscriptions(me) if a.name == agent_name
    )

    def auth() -> dict:
        return {"Authorization": f"Bearer {runtime.token_manager.get_access_token()}"}

    async with httpx.AsyncClient(timeout=300.0) as http:
        # The subscription gives you a *relay* card URL. Fetch it and read the
        # JSON-RPC endpoint out of the card rather than guessing a path.
        card = (await http.get(str(target.agent_card_url), headers=auth())).json()
        rpc_url = card["url"]

        send = await http.post(rpc_url, headers=auth(), json={
            "jsonrpc": "2.0",
            "id": 1,
            "method": "message/send",
            "params": {"message": {
                "role": "user",
                "parts": [{"kind": "text", "text": prompt}],
                "messageId": str(uuid4()),
            }},
        })
        task = send.json()["result"]

        # This is the bit the wrappers exist for: A2A tasks come back
        # non-terminal, and an LLM handed {"state": "working"} is useless.
        while task.get("status", {}).get("state") not in TERMINAL:
            await asyncio.sleep(5)
            poll = await http.post(rpc_url, headers=auth(), json={
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tasks/get",
                "params": {"id": task["id"]},
            })
            task = poll.json()["result"]

        return task


print(asyncio.run(ask("weather_agent", "What's the weather in Tokyo?")))