"""
Class for selecting an agent to execute an instruction in the multi-agent workflow.
"""

from abc import ABC, abstractmethod
from itertools import cycle
from typing import Dict, List, Optional

from loguru import logger
from pydantic import BaseModel

from elemental_agents.core.agent.agent import Agent
from elemental_agents.utils.exceptions import AgentException
from elemental_agents.utils.utils import extract_tag_content


class AgentSelectorParameters(BaseModel):
    """
    Parameters for the agent selector.
    """

    agents: Dict[str, Agent]
    last_message: str | None = None
    execution_order: Optional[List[str]] = None

    class Config:
        """
        Allow for Agent to be used as a key in the dictionary.
        """

        arbitrary_types_allowed = True


class AgentSelector(ABC):
    """
    Abstract class for selecting an agent to execute an instruction in the
    multi-agent workflow. The selection is to be done for the next agent to
    execute the instruction based on the instruction itself or the previous
    agent's response. The selection can be done in an iterative, intelligent, or
    conversational manner.
    """

    def __init__(self, lead_agent: str) -> None:
        """
        Initialize the agent selector with the lead agent.

        :param lead_agent: The lead agent for the conversation.
        """

        self._lead_agent: str = lead_agent
        self._current_agent: str = None

    @abstractmethod
    def select(self, parameters: AgentSelectorParameters) -> Agent:
        """
        Select an agent to execute the instruction based on the instruction itself.

        :param parameters: The parameters for the agent selector.
        :return: The selected agent executor.
        """


class ConversationalSelector(AgentSelector):
    """
    Selects the agent based on the last message received and the <next> section
    of the message generated by the agent. The @AgentName is extracted from the
    <next> section and the agent with the name is selected.
    """

    def __init__(self, lead_agent: str) -> None:
        super().__init__(lead_agent)
        logger.info(f"Conversational selector initialized. Lead agent: {lead_agent}")

    def select(self, parameters: AgentSelectorParameters) -> Agent:
        """
        Select an agent to execute the instruction based on the last message
        received and <next> section with @AgentName tag.

        :param parameters: The parameters for the agent selector.
        :return: The selected agent executor.
        """

        # Initial message, no previous message present
        if parameters.last_message is None:
            self._current_agent = self._lead_agent
            logger.debug("No previous message present. Selecting the lead agent.")
            return parameters.agents[self._lead_agent]

        try:
            next_agents = extract_tag_content(parameters.last_message, "next")
            next_agent = next_agents[0].split("@")[-1].strip()

            if (next_agent == "Self") and (self._current_agent is not None):
                next_agent = self._current_agent

            # Check if 'next_agent' is in the list of agents
            if next_agent not in parameters.agents:
                logger.error(
                    (
                        f"Agent {next_agent} is not registered with the agent unit. "
                        "Proceeding with the lead agent."
                    )
                )
                next_agent = self._lead_agent

            self._current_agent = next_agent
            logger.debug(f"Next agent: {next_agent}")
        except IndexError:
            logger.error(
                "No next agent found in the message. Selecting the lead agent."
            )
            self._current_agent = self._lead_agent
            return parameters.agents[self._lead_agent]

        return parameters.agents[next_agent]


class IdentitySelector(AgentSelector):
    """
    This selector always selects the lead agent.
    """

    def __init__(self, lead_agent: str) -> None:
        super().__init__(lead_agent)
        logger.info(f"Identity selector initialized. Lead agent: {lead_agent}")

    def select(self, parameters: AgentSelectorParameters) -> Agent:
        """
        Identity selector selects always the lead agent.

        :param parameters: The parameters for the agent selector.
        :return: The selected agent executor.
        """
        self._current_agent = self._lead_agent

        return parameters.agents[self._lead_agent]


class IterativeSelector(AgentSelector):
    """
    Select agent based of predefined order and always cycle through the provided
    list of names to select next agent.
    """

    def __init__(self, lead_agent: str) -> None:
        super().__init__(lead_agent)
        self._cyclic_iterator: cycle | None = None
        logger.info(f"Iterative selector initialized. Lead agent: {lead_agent}")

    def select(self, parameters: AgentSelectorParameters) -> Agent:
        """
        Select next agent based of the list provided in the workflow
        configuration and cycle through the elements of this list.

        :param parameters: The parameters for the agent selector.
        :return: The selected agent executor.
        """
        if parameters.execution_order is None:
            raise AgentException(
                "Execution order not provided in the agent selector parameters."
            )

        if self._cyclic_iterator is None:
            self._cyclic_iterator = cycle(parameters.execution_order)

        next_agent = next(self._cyclic_iterator)
        logger.debug(f"Next agent: {next_agent}")
        self._current_agent = next_agent

        if next_agent not in parameters.agents:
            raise AgentException(
                f"Agent {next_agent} not registered with the agent unit."
            )
        return parameters.agents[next_agent]
