# Copyright (c) Microsoft. All rights reserved.


import asyncio

from agent_framework import Agent, Message
from dotenv import load_dotenv

"""AutoGen SelectorGroupChat vs Agent Framework GroupChatBuilder.

Demonstrates LLM-based speaker selection where an orchestrator decides
which agent should speak next based on the conversation context.
"""

# Load environment variables from .env file
load_dotenv()


async def run_autogen() -> None:
    """AutoGen's SelectorGroupChat with LLM-based speaker selection."""

    from autogen_agentchat.agents import AssistantAgent
    from autogen_agentchat.conditions import MaxMessageTermination
    from autogen_agentchat.teams import SelectorGroupChat
    from autogen_agentchat.ui import Console
    from autogen_ext.models.openai import OpenAIChatCompletionClient

    client = OpenAIChatCompletionClient(model="gpt-4.1-mini")

    # Create specialized agents
    python_expert = AssistantAgent(
        name="python_expert",
        model_client=client,
        system_message="You are a Python programming expert. Answer Python-related questions.",
        description="Expert in Python programming",
        model_client_stream=True,
    )

    javascript_expert = AssistantAgent(
        name="javascript_expert",
        model_client=client,
        system_message="You are a JavaScript programming expert. Answer JavaScript-related questions.",
        description="Expert in JavaScript programming",
        model_client_stream=True,
    )

    database_expert = AssistantAgent(
        name="database_expert",
        model_client=client,
        system_message="You are a database expert. Answer SQL and database-related questions.",
        description="Expert in databases and SQL",
        model_client_stream=True,
    )

    # Create selector group chat - LLM selects appropriate expert
    team = SelectorGroupChat(
        participants=[python_expert, javascript_expert, database_expert],
        model_client=client,
        termination_condition=MaxMessageTermination(2),
        selector_prompt="Based on the conversation so far:\n{history}\n, "
        "select the most appropriate expert from {roles} to respond next.",
    )

    # Run with a question that requires expert selection
    print("[AutoGen] Selector group chat conversation:")
    await Console(team.run_stream(task="How do I connect to a PostgreSQL database using Python?"))


async def run_agent_framework() -> None:
    """Agent Framework's GroupChatBuilder with LLM-based speaker selection."""
    from agent_framework.openai import OpenAIChatClient
    from agent_framework.orchestrations import GroupChatBuilder

    client = OpenAIChatClient(model="gpt-4.1-mini")

    # Create specialized agents
    python_expert = Agent(
        client=client,
        name="python_expert",
        instructions="You are a Python programming expert. Answer Python-related questions.",
        description="Expert in Python programming",
    )

    javascript_expert = Agent(
        client=client,
        name="javascript_expert",
        instructions="You are a JavaScript programming expert. Answer JavaScript-related questions.",
        description="Expert in JavaScript programming",
    )

    database_expert = Agent(
        client=client,
        name="database_expert",
        instructions="You are a database expert. Answer SQL and database-related questions.",
        description="Expert in databases and SQL",
    )

    workflow = GroupChatBuilder(
        participants=[python_expert, javascript_expert, database_expert],
        max_rounds=1,
        orchestrator_agent=Agent(
            client=client,
            name="selector_manager",
            instructions="Based on the conversation, select the most appropriate expert to respond next.",
        ),
    ).build()

    # Run with a question that requires expert selection
    print("[Agent Framework] Group chat conversation:")
    async for event in workflow.run("How do I connect to a PostgreSQL database using Python?", stream=True):
        if event.type == "output" and isinstance(event.data, list):
            for message in event.data:  # type: ignore
                if isinstance(message, Message) and message.role == "assistant" and message.text:
                    print(f"---------- {message.author_name} ----------")
                    print(message.text)


async def main() -> None:
    print("=" * 60)
    print("Selector Group Chat Comparison")
    print("=" * 60)
    print("AutoGen: SelectorGroupChat")
    print("Agent Framework: GroupChatBuilder with standard_manager\n")
    await run_autogen()
    print()
    await run_agent_framework()


if __name__ == "__main__":
    asyncio.run(main())
