from typing import Any, Callable, Dict, Optional, Sequence

from swarms.structs.base_swarm import BaseSwarm
from swarms.utils.loguru_logger import logger


class AutoSwarmRouter(BaseSwarm):
    """AutoSwarmRouter class represents a router for the AutoSwarm class.

    This class is responsible for routing tasks to the appropriate swarm based on the provided name.
    It allows customization of the preprocessing, routing, and postprocessing of tasks.

    Attributes:
        name (str): The name of the router.
        description (str): The description of the router.
        verbose (bool): Whether to enable verbose mode.
        custom_params (dict): Custom parameters for the router.
        swarms (list): A list of BaseSwarm objects.
        custom_preprocess (callable): Custom preprocessing function for tasks.
        custom_postprocess (callable): Custom postprocessing function for task results.
        custom_router (callable): Custom routing function for tasks.

    Methods:
        run(task: str = None, *args, **kwargs) -> Any:
            Run the swarm simulation and route the task to the appropriate swarm.

    Flow:
        name -> router -> swarm entry point
    """

    def __init__(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
        verbose: bool = False,
        custom_params: Optional[Dict[str, Any]] = None,
        swarms: Sequence[BaseSwarm] = None,
        custom_preprocess: Optional[Callable] = None,
        custom_postprocess: Optional[Callable] = None,
        custom_router: Optional[Callable] = None,
        *args,
        **kwargs,
    ):
        super().__init__(
            name=name, description=description, *args, **kwargs
        )
        self.name = name
        self.description = description
        self.verbose = verbose
        self.custom_params = custom_params
        self.swarms = swarms
        self.custom_preprocess = custom_preprocess
        self.custom_postprocess = custom_postprocess
        self.custom_router = custom_router

        # Create a dictionary of swarms
        self.swarm_dict = {swarm.name: swarm for swarm in self.swarms}

        logger.info(
            f"AutoSwarmRouter has been initialized with {self.len_of_swarms()} swarms."
        )

    def run(self, task: str = None, *args, **kwargs):
        try:
            """Run the swarm simulation and route the task to the appropriate swarm."""

            if self.custom_preprocess:
                # If custom preprocess function is provided then run it
                logger.info("Running custom preprocess function.")
                task, args, kwargs = self.custom_preprocess(
                    task, args, kwargs
                )

            if self.custom_router:
                # If custom router function is provided then use it to route the task
                logger.info("Running custom router function.")
                out = self.custom_router(self, task, *args, **kwargs)

                if self.custom_postprocess:
                    # If custom postprocess function is provided then run it
                    out = self.custom_postprocess(out)

                return out

            if self.name in self.swarm_dict:
                # If a match is found then send the task to the swarm
                out = self.swarm_dict[self.name].run(
                    task, *args, **kwargs
                )

                if self.custom_postprocess:
                    # If custom postprocess function is provided then run it
                    out = self.custom_postprocess(out)

                return out

            # If no match is found then return None
            raise ValueError(
                f"Swarm with name {self.name} not found."
            )
        except Exception as e:
            logger.error(f"Error: {e}")
            raise e

    def len_of_swarms(self):
        return print(len(self.swarms))

    def list_available_swarms(self):
        for swarm in self.swarms:
            try:
                logger.info(
                    f"Swarm Name: {swarm.name} || Swarm Description: {swarm.description} "
                )
            except Exception as error:
                logger.error(
                    f"Error Detected You may not have swarms available: {error}"
                )
                raise error


class AutoSwarm(BaseSwarm):
    """AutoSwarm class represents a swarm of agents that can be created automatically.

    Flow:
    name -> router -> swarm entry point

    Args:
        name (Optional[str]): The name of the swarm. Defaults to None.
        description (Optional[str]): The description of the swarm. Defaults to None.
        verbose (bool): Whether to enable verbose mode. Defaults to False.
        custom_params (Optional[Dict[str, Any]]): Custom parameters for the swarm. Defaults to None.
        router (Optional[AutoSwarmRouter]): The router for the swarm. Defaults to None.
    """

    def __init__(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
        verbose: bool = False,
        custom_params: Optional[Dict[str, Any]] = None,
        custom_preprocess: Optional[Callable] = None,
        custom_postprocess: Optional[Callable] = None,
        custom_router: Optional[Callable] = None,
        max_loops: int = 1,
        *args,
        **kwargs,
    ):
        super().__init__()
        self.name = name
        self.description = description
        self.verbose = verbose
        self.custom_params = custom_params
        self.custom_preprocess = custom_preprocess
        self.custom_postprocess = custom_postprocess
        self.custom_router = custom_router
        self.max_loops = max_loops
        self.router = AutoSwarmRouter(
            name=name,
            description=description,
            verbose=verbose,
            custom_params=custom_params,
            custom_preprocess=custom_preprocess,
            custom_postprocess=custom_postprocess,
            custom_router=custom_router,
            *args,
            **kwargs,
        )

        if name is None:
            raise ValueError(
                "A name must be provided for the AutoSwarm, what swarm do you want to use?"
            )

        if verbose is True:
            self.init_logging()

    def init_logging(self):
        logger.info("AutoSwarm has been activated. Ready for usage.")

    # def name_swarm_check(self, name: str = None):

    def run(self, task: str = None, *args, **kwargs):
        """Run the swarm simulation."""
        try:
            loop = 0

            while loop < self.max_loops:
                if self.custom_preprocess:
                    # If custom preprocess function is provided then run it
                    logger.info("Running custom preprocess function.")
                    task, args, kwargs = self.custom_preprocess(
                        task, args, kwargs
                    )

                if self.custom_router:
                    # If custom router function is provided then use it to route the task
                    logger.info("Running custom router function.")
                    out = self.custom_router(
                        self, task, *args, **kwargs
                    )

                else:
                    out = self.router.run(task, *args, **kwargs)

                if self.custom_postprocess:
                    # If custom postprocess function is provided then run it
                    out = self.custom_postprocess(out)

                # LOOP
                loop += 1

                return out
        except Exception as e:
            logger.error(
                f"Error: {e} try optimizing the inputs and try again."
            )
            raise e

    def list_all_swarms(self):
        for swarm in self.swarms:
            try:
                logger.info(
                    f"Swarm Name: {swarm.name} || Swarm Description: {swarm.description} "
                )
            except Exception as error:
                logger.error(
                    f"Error Detected You may not have swarms available: {error}"
                )
                raise error