import asyncio
from .consensus import ConsensusLayer, OutputAggregator


class AgentSwarm:
    def __init__(self, agent_list):
        """
        Initialize the AgentSwarm with a list of agents, a consensus layer, and an output aggregator.

        Parameters
        ----------
        agent_list : list
            A list of agent objects that will be used to generate responses and participate in consensus.
        """
        self.agent_list = agent_list
        self.consensus_layer = ConsensusLayer()
        self.aggregator = OutputAggregator()

    async def _gather_responses(self, agent_list, query):
        """
        Gather responses from all agents asynchronously.

        Parameters
        ----------
        agent_list : list
            A list of agent objects that will generate responses.
        query : str
            The query or request to be sent to the agents for generating responses.

        Returns
        -------
        list
            A list of responses generated by each agent in the agent_list.
        """
        return await asyncio.gather(*(a.generate_response(query) for a in agent_list))

    def run_iteration(self, query):
        """
        Run one iteration of the AgentSwarm, which gathers responses, computes consensus, and aggregates the output.

        Parameters
        ----------
        query : str
            The query to be processed by the agents in the swarm.

        Returns
        -------
        object
            The aggregated output based on the consensus of the agent responses.
        """
        # Get the consensus
        responses = asyncio.run(self._gather_responses(self.agent_list, query))
        consensus_result = self.consensus_layer.get_consensus(
            responses, self.agent_list)
        output = self.aggregator.aggregate_output(consensus_result)

        return output
