import concurrent.futures
import json
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Union

from termcolor import colored

from swarms.structs.agent import Agent
from swarms.structs.task import Task
from swarms.utils.logger import logger


# SequentialWorkflow class definition using dataclasses
@dataclass
class SequentialWorkflow:
    """
    SequentialWorkflow class for running a sequence of tasks using N number of autonomous agents.

    Args:
        max_loops (int): The maximum number of times to run the workflow.
        dashboard (bool): Whether to display the dashboard for the workflow.


    Attributes:
        tasks (List[Task]): The list of tasks to execute.
        max_loops (int): The maximum number of times to run the workflow.
        dashboard (bool): Whether to display the dashboard for the workflow.


    Examples:
    >>> from swarms.models import OpenAIChat
    >>> from swarms.structs import SequentialWorkflow
    >>> llm = OpenAIChat(openai_api_key="")
    >>> workflow = SequentialWorkflow(max_loops=1)
    >>> workflow.add("What's the weather in miami", llm)
    >>> workflow.add("Create a report on these metrics", llm)
    >>> workflow.run()
    >>> workflow.tasks

    """

    name: str = None
    description: str = None
    tasks: List[Task] = field(default_factory=list)
    max_loops: int = 1
    autosave: bool = False
    saved_state_filepath: Optional[str] = (
        "sequential_workflow_state.json"
    )
    restore_state_filepath: Optional[str] = None
    dashboard: bool = False

    def add(
        self,
        agent: Union[Callable, Agent],
        task: Optional[str] = None,
        tasks: Optional[List[str]] = None,
        *args,
        **kwargs,
    ) -> None:
        """
        Add a task to the workflow.

        Args:
            agent (Union[Callable, Agent]): The model or agent to execute the task.
            task (str): The task description or the initial input for the Agent.

            *args: Additional arguments to pass to the task execution.
            **kwargs: Additional keyword arguments to pass to the task execution.
        """
        try:
            # If the agent is a Agent instance, we include the task in kwargs for Agent.run()
            if isinstance(agent, Agent):
                kwargs["task"] = (
                    task  # Set the task as a keyword argument for Agent
                )

            # Append the task to the tasks list
            self.tasks.append(
                Task(
                    description=task,
                    agent=agent,
                    args=list(args),
                    kwargs=kwargs,
                )
            )

            logger.info(
                f"[INFO][SequentialWorkflow] Added task {task} to"
                " workflow"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error adding task to workflow: {error}", "red"
                ),
            )

    def reset_workflow(self) -> None:
        """Resets the workflow by clearing the results of each task."""
        try:
            for task in self.tasks:
                task.result = None
        except Exception as error:
            logger.error(
                colored(f"Error resetting workflow: {error}", "red"),
            )

    def get_task_results(self) -> Dict[str, Any]:
        """
        Returns the results of each task in the workflow.

        Returns:
            Dict[str, Any]: The results of each task in the workflow
        """
        try:
            return {
                task.description: task.result for task in self.tasks
            }
        except Exception as error:
            logger.error(
                colored(
                    f"Error getting task results: {error}", "red"
                ),
            )

    def remove_task(self, task: str) -> None:
        """Remove tasks from sequential workflow"""
        try:
            self.tasks = [
                task
                for task in self.tasks
                if task.description != task
            ]
            logger.info(
                f"[INFO][SequentialWorkflow] Removed task {task} from"
                " workflow"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error removing task from workflow: {error}",
                    "red",
                ),
            )

    def update_task(self, task: str, **updates) -> None:
        """
        Updates the arguments of a task in the workflow.

        Args:
            task (str): The description of the task to update.
            **updates: The updates to apply to the task.

        Raises:
            ValueError: If the task is not found in the workflow.

        Examples:
        >>> from swarms.models import OpenAIChat
        >>> from swarms.structs import SequentialWorkflow
        >>> llm = OpenAIChat(openai_api_key="")
        >>> workflow = SequentialWorkflow(max_loops=1)
        >>> workflow.add("What's the weather in miami", llm)
        >>> workflow.add("Create a report on these metrics", llm)
        >>> workflow.update_task("What's the weather in miami", max_tokens=1000)
        >>> workflow.tasks[0].kwargs
        {'max_tokens': 1000}

        """
        try:
            for task in self.tasks:
                if task.description == task:
                    task.kwargs.update(updates)
                    break
            else:
                raise ValueError(
                    f"Task {task} not found in workflow."
                )

            print(
                f"[INFO][SequentialWorkflow] Updated task {task} in"
                " workflow"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error updating task in workflow: {error}", "red"
                ),
            )

    def delete_task(self, task: str) -> None:
        """
        Delete a task from the workflow.

        Args:
            task (str): The description of the task to delete.

        Raises:
            ValueError: If the task is not found in the workflow.

        Examples:
        >>> from swarms.models import OpenAIChat
        >>> from swarms.structs import SequentialWorkflow
        >>> llm = OpenAIChat(openai_api_key="")
        >>> workflow = SequentialWorkflow(max_loops=1)
        >>> workflow.add("What's the weather in miami", llm)
        >>> workflow.add("Create a report on these metrics", llm)
        >>> workflow.delete_task("What's the weather in miami")
        >>> workflow.tasks
        [Task(description='Create a report on these metrics', agent=Agent(llm=OpenAIChat(openai_api_key=''), max_loops=1, dashboard=False), args=[], kwargs={}, result=None, history=[])]
        """
        try:
            for task in self.tasks:
                if task.description == task:
                    self.tasks.remove(task)
                    break
            else:
                raise ValueError(
                    f"Task {task} not found in workflow."
                )

            print(
                f"[INFO][SequentialWorkflow] Deleted task {task} from"
                " workflow"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error deleting task from workflow: {error}",
                    "red",
                ),
            )

    def concurrent_run(self):
        """
        Concurrently run the workflow using a pool of workers.

        Examples:
        >>> from swarms.models import OpenAIChat
        >>> from swarms.structs import SequentialWorkflow
        >>> llm = OpenAIChat(openai_api_key="")
        >>> workflow = SequentialWorkflow(max_loops=1)

        """
        try:
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures_to_task = {
                    executor.submit(task.run): task
                    for task in self.tasks
                }
                results = []
                for future in concurrent.futures.as_completed(
                    futures_to_task
                ):
                    task = futures_to_task[future]

                    try:
                        result = future.result()
                    except Exception as error:
                        print(f"Error running workflow: {error}")
                    else:
                        results.append(result)
                        print(
                            f"Task {task} completed successfully with"
                            f" result: {result}"
                        )
        except Exception as error:
            print(colored(f"Error running workflow: {error}", "red"))

    def save_workflow_state(
        self,
        filepath: Optional[str] = "sequential_workflow_state.json",
        **kwargs,
    ) -> None:
        """
        Saves the workflow state to a json file.

        Args:
            filepath (str): The path to save the workflow state to.

        Examples:
        >>> from swarms.models import OpenAIChat
        >>> from swarms.structs import SequentialWorkflow
        >>> llm = OpenAIChat(openai_api_key="")
        >>> workflow = SequentialWorkflow(max_loops=1)
        >>> workflow.add("What's the weather in miami", llm)
        >>> workflow.add("Create a report on these metrics", llm)
        >>> workflow.save_workflow_state("sequential_workflow_state.json")
        """
        try:
            filepath = filepath or self.saved_state_filepath

            with open(filepath, "w") as f:
                # Saving the state as a json for simplicuty
                state = {
                    "tasks": [
                        {
                            "description": task.description,
                            "args": task.args,
                            "kwargs": task.kwargs,
                            "result": task.result,
                            "history": task.history,
                        }
                        for task in self.tasks
                    ],
                    "max_loops": self.max_loops,
                }
                json.dump(state, f, indent=4)

            print(
                "[INFO][SequentialWorkflow] Saved workflow state to"
                f" {filepath}"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error saving workflow state: {error}",
                    "red",
                )
            )

    def workflow_bootup(self, **kwargs) -> None:
        """
        Workflow bootup.

        """
        print(
            colored(
                """
                Sequential Workflow Initializing...""",
                "green",
                attrs=["bold", "underline"],
            )
        )

    def workflow_dashboard(self, **kwargs) -> None:
        """
        Displays a dashboard for the workflow.

        Args:
            **kwargs: Additional keyword arguments to pass to the dashboard.

        Examples:
        >>> from swarms.models import OpenAIChat
        >>> from swarms.structs import SequentialWorkflow
        >>> llm = OpenAIChat(openai_api_key="")
        >>> workflow = SequentialWorkflow(max_loops=1)
        >>> workflow.add("What's the weather in miami", llm)
        >>> workflow.add("Create a report on these metrics", llm)
        >>> workflow.workflow_dashboard()

        """
        print(
            colored(
                f"""
                Sequential Workflow Dashboard
                --------------------------------
                Name: {self.name}
                Description: {self.description}
                Tasks: {len(self.tasks)}
                Max Loops: {self.max_loops}
                Autosave: {self.autosave}
                Autosave Filepath: {self.saved_state_filepath}
                Restore Filepath: {self.restore_state_filepath}
                --------------------------------
                Metadata:
                kwargs: {kwargs}
                """,
                "cyan",
                attrs=["bold", "underline"],
            )
        )

    def workflow_shutdown(self, **kwargs) -> None:
        """Shuts down the workflow."""
        print(
            colored(
                """
                Sequential Workflow Shutdown...""",
                "red",
                attrs=["bold", "underline"],
            )
        )

    def add_objective_to_workflow(self, task: str, **kwargs) -> None:
        """Adds an objective to the workflow."""
        try:
            print(
                colored(
                    """
                    Adding Objective to Workflow...""",
                    "green",
                    attrs=["bold", "underline"],
                )
            )

            task = Task(
                description=task,
                agent=kwargs["agent"],
                args=list(kwargs["args"]),
                kwargs=kwargs["kwargs"],
            )
            self.tasks.append(task)

            print(
                f"[INFO][SequentialWorkflow] Added task {task} to"
                " workflow"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error adding objective to workflow: {error}",
                    "red",
                )
            )

    def load_workflow_state(
        self, filepath: str = None, **kwargs
    ) -> None:
        """
        Loads the workflow state from a json file and restores the workflow state.

        Args:
            filepath (str): The path to load the workflow state from.

        Examples:
        >>> from swarms.models import OpenAIChat
        >>> from swarms.structs import SequentialWorkflow
        >>> llm = OpenAIChat(openai_api_key="")
        >>> workflow = SequentialWorkflow(max_loops=1)
        >>> workflow.add("What's the weather in miami", llm)
        >>> workflow.add("Create a report on these metrics", llm)
        >>> workflow.save_workflow_state("sequential_workflow_state.json")
        >>> workflow.load_workflow_state("sequential_workflow_state.json")

        """
        try:
            filepath = filepath or self.restore_state_filepath

            with open(filepath, "r") as f:
                state = json.load(f)
                self.max_loops = state["max_loops"]
                self.tasks = []
                for task_state in state["tasks"]:
                    task = Task(
                        description=task_state["description"],
                        agent=task_state["agent"],
                        args=task_state["args"],
                        kwargs=task_state["kwargs"],
                        result=task_state["result"],
                        history=task_state["history"],
                    )
                    self.tasks.append(task)

            print(
                "[INFO][SequentialWorkflow] Loaded workflow state"
                f" from {filepath}"
            )
        except Exception as error:
            logger.error(
                colored(
                    f"Error loading workflow state: {error}",
                    "red",
                )
            )

    def run(self) -> None:
        """
        Run the workflow.

        Raises:
            ValueError: If a Agent instance is used as a task and the 'task' argument is not provided.

        """
        try:
            self.workflow_bootup()
            for _ in range(self.max_loops):
                for task in self.tasks:
                    # Check if the current task can be executed
                    if task.result is None:
                        # Check if the agent is a Agent and a 'task' argument is needed
                        if isinstance(task.agent, Agent):
                            # Ensure that 'task' is provided in the kwargs
                            if "task" not in task.kwargs:
                                raise ValueError(
                                    "The 'task' argument is required"
                                    " for the Agent agent execution"
                                    f" in '{task.description}'"
                                )
                            # Separate the 'task' argument from other kwargs
                            flow_task_arg = task.kwargs.pop("task")
                            task.result = task.agent.run(
                                flow_task_arg,
                                *task.args,
                                **task.kwargs,
                            )
                        else:
                            # If it's not a Agent instance, call the agent directly
                            task.result = task.agent(
                                *task.args, **task.kwargs
                            )

                        # Pass the result as an argument to the next task if it exists
                        next_task_index = self.tasks.index(task) + 1
                        if next_task_index < len(self.tasks):
                            next_task = self.tasks[next_task_index]
                            if isinstance(next_task.agent, Agent):
                                # For Agent flows, 'task' should be a keyword argument
                                next_task.kwargs["task"] = task.result
                            else:
                                # For other callable flows, the result is added to args
                                next_task.args.insert(0, task.result)

                        # Autosave the workflow state
                        if self.autosave:
                            self.save_workflow_state(
                                "sequential_workflow_state.json"
                            )
        except Exception as e:
            logger.error(
                colored(
                    (
                        "Error initializing the Sequential workflow:"
                        f" {e} try optimizing your inputs like the"
                        " agent class and task description"
                    ),
                    "red",
                    attrs=["bold", "underline"],
                )
            )

    async def arun(self) -> None:
        """
        Asynchronously run the workflow.

        Raises:
            ValueError: If a Agent instance is used as a task and the 'task' argument is not provided.

        """
        try:
            for _ in range(self.max_loops):
                for task in self.tasks:
                    # Check if the current task can be executed
                    if task.result is None:
                        # Check if the agent is a Agent and a 'task' argument is needed
                        if isinstance(task.agent, Agent):
                            # Ensure that 'task' is provided in the kwargs
                            if "task" not in task.kwargs:
                                raise ValueError(
                                    "The 'task' argument is required"
                                    " for the Agent agent execution"
                                    f" in '{task.description}'"
                                )
                            # Separate the 'task' argument from other kwargs
                            flow_task_arg = task.kwargs.pop("task")
                            task.result = await task.agent.arun(
                                flow_task_arg,
                                *task.args,
                                **task.kwargs,
                            )
                        else:
                            # If it's not a Agent instance, call the agent directly
                            task.result = await task.agent(
                                *task.args, **task.kwargs
                            )

                        # Pass the result as an argument to the next task if it exists
                        next_task_index = self.tasks.index(task) + 1
                        if next_task_index < len(self.tasks):
                            next_task = self.tasks[next_task_index]
                            if isinstance(next_task.agent, Agent):
                                # For Agent flows, 'task' should be a keyword argument
                                next_task.kwargs["task"] = task.result
                            else:
                                # For other callable flows, the result is added to args
                                next_task.args.insert(0, task.result)

                        # Autosave the workflow state
                        if self.autosave:
                            self.save_workflow_state(
                                "sequential_workflow_state.json"
                            )
        except Exception as e:
            print(
                colored(
                    (
                        "Error initializing the Sequential workflow:"
                        f" {e} try optimizing your inputs like the"
                        " agent class and task description"
                    ),
                    "red",
                    attrs=["bold", "underline"],
                )
            )
