# import asyncio
# from actual.society_of_minds.society_of_minds import create_swarm_mind
# from autogen_agentchat.agents import AssistantAgent
# from autogen_agentchat.teams import SelectorGroupChat
# from autogen_agentchat.conditions import TextMentionTermination
# from autogen_agentchat.ui import Console
# from neuroengine.client import NeuroengineChatCompletionClient, neuro_client_factory
# from actual.reconnaissance.nmap import nmap_society_of_mind
# from actual.reconnaissance.web_directory_enumeration import web_directory_enumeration_society_of_mind

# planner_agent_name: str = "ReconnaissancePlannerAgent"
# planner_neuro_client: NeuroengineChatCompletionClient = neuro_client_factory(name=planner_agent_name)
# planner_system_message: str = """
#                                You are a cybersecurity expert, and you are assigned to reconnaissance planning coordinator.
#                                Your job is to break down complex reconnaissance task into smaller, manageable subtasks."
#                                Coordinate reconnaissance task by delegating to specialized agents:"
#                                - nmap_society_of_mind: For scanning of ports"
#                                - web_directory_enumeration_society_of_mind: For enumeration of a web server to determine the subdirectories and files"

#                                You only plan and delegate task - you do not execute them yourself.

#                                When assigning tasks, use this format:
#                                1. <agent> : <task>

#                                Once the reconnaissance task is finished, summarize findings and say TERMINATE when it ends.
#                                """

# planner_description: str = "This agent plans and delegates tasks to other agents."

# ReconnaissancePlannerAgent: AssistantAgent = AssistantAgent(name=planner_agent_name, model_client=planner_neuro_client, system_message=planner_system_message, description=planner_description, model_client_stream=True)


# selector_prompt: str = """

# Based on the agent descriptions and current context, select the most appropriate agent to handle the task.

# {roles}

# Current conversation context:
# {history}

# Read the above conversation, then select an agent from {participants} to perform the next task.
# Make sure the planner agent has assigned tasks before other agents start working.
# Only select one agent.
# """

# terminate: TextMentionTermination = TextMentionTermination("TERMINATE")

# recon_neuro_client_name: str = "Reconnaissance"
# recon_neuro_client: NeuroengineChatCompletionClient = neuro_client_factory(name=recon_neuro_client_name)


# team: SelectorGroupChat = SelectorGroupChat(participants=[ReconnaissancePlannerAgent, nmap_society_of_mind, web_directory_enumeration_society_of_mind], model_client=recon_neuro_client, selector_prompt=selector_prompt, termination_condition=terminate)

# task: str = "Do reconnaissance and gather information on localhost"


# async def main():
#     await Console(team.run_stream(task=task))

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

import asyncio
from neuroengine.client import NeuroengineChatCompletionClient, neuro_client_factory
from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
from autogen_agentchat.teams import SelectorGroupChat
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.ui import Console


async def main() -> None:

    async def lookup_hotel(location: str) -> str:
        """
        Helps to locate suitable hotels at a given location.
        """
        return f"Here are some hotels in {location}: hotel1, hotel2, hotel3."

    async def lookup_flight(origin: str, destination: str) -> str:
        """
        Helps lookup flights from origin to destination and book them.
        """
        return f"Here are some flights from {origin} to {destination}: flight1, flight2, flight3."

    async def book_trip() -> str:
        """
        Helps book and confirm the trip
        """
        return "Your trip is booked!"

    # travel_advisor = AssistantAgent(
    #     name="Travel_Planner",
    #     model_client=neuro_client_factory(name="Travel_Planner"),
    #     tools=[book_trip],
    #     description="Plans what are the tasks need to be done. Returns a numbered list of agents and at what order they need to be executed.",
    #     system_message="Plans what are the tasks need to be done. Returns a numbered list of agents and at what order they need to be executed. You have Hotel_Agent and Flight_Agent to help you.",
    #     model_client_stream=True
    # )

    hotel_agent = AssistantAgent(
        name="Hotel_Agent",
        model_client=neuro_client_factory(name="Hotel_Agent"),
        tools=[lookup_hotel],
        description="Helps with locating hotels for trip and confirming the hotels booking.",
        system_message="Helps with locating hotels for trip and confirming the hotels booking.",
        model_client_stream=True

    )

    flight_agent = AssistantAgent(
        name="Flight_Agent",
        model_client=neuro_client_factory(name="Flight_Agent"),
        tools=[lookup_flight],
        description="Helps with locating flights for trip and confirming the flights booking.",
        system_message="Helps with locating flights for trip and confirming the flights booking.",
        model_client_stream=True

    )

    user_proxy = UserProxyAgent(name="user_agent", description="Used for terminating the task when it is completed.")

    selector_prompt: str = """
    You have the following roles:
    {roles}

    Current conversation history:
    {history}

    Based on the history, select from {participants} the best agent suited for the role
    """
    termination = TextMentionTermination("TERMINATE")
    team = SelectorGroupChat(
        [ hotel_agent, flight_agent, user_proxy],
        model_client=neuro_client_factory(name="selector"),
        termination_condition=termination,
        allow_repeated_speaker=False,
        model_client_streaming=True
    )

    await Console(team.run_stream(task="Book a 3-day trip to new york."))


asyncio.run(main())