import os

from autogen import AfterWork, AfterWorkOption, ConversableAgent, OnCondition, initiate_swarm_chat, register_hand_off
from autogen.agents.experimental import DiscordAgent, WebSurferAgent

# LLM configuration for our agent to select the tools and craft the message
llm_config = {"model": "gpt-4o-mini"}

triage_agent = ConversableAgent(
    name="triage_agent",
    llm_config=llm_config,
)

websurfer_agent = WebSurferAgent(
    name="websurfer_agent",
    llm_config=llm_config,
)

discord_agent = DiscordAgent(
    name="discord_agent",
    llm_config=llm_config,
    bot_token=os.getenv("DISCORD_BOT_TOKEN"),
    guild_name="My Discord Server",
    channel_name="general",
    is_termination_msg=lambda x: (x["content"] is not None) and "TERMINATE" in x["content"],
)

register_hand_off(
    agent=triage_agent,
    hand_to=[
        OnCondition(target=websurfer_agent, condition="Find information on the web."),
        OnCondition(
            target=discord_agent,
            condition="Send weather details to Discord only after you have received them.",
        ),
        AfterWork(AfterWorkOption.TERMINATE),
    ],
)

register_hand_off(agent=websurfer_agent, hand_to=[AfterWork(agent=triage_agent)])

register_hand_off(agent=discord_agent, hand_to=[AfterWork(agent=triage_agent)])

result, cv, last_agent = initiate_swarm_chat(
    initial_agent=triage_agent,
    agents=[triage_agent, discord_agent, websurfer_agent],
    messages="Find the latest weather for Sydney, NSW, Australia and post it on Discord",
    max_rounds=10,
)