import os
import uuid
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Union

from pydantic import BaseModel, Field

from swarms.structs.agent import Agent
from swarms.structs.base_swarm import BaseSwarm
from swarms.utils.file_processing import create_file_in_folder
from swarms.utils.loguru_logger import initialize_logger
from swarms.structs.conversation import Conversation
from swarms.structs.swarm_id_generator import generate_swarm_id
from swarms.structs.output_types import OutputType

logger = initialize_logger(log_folder="concurrent_workflow")


class AgentOutputSchema(BaseModel):
    run_id: Optional[str] = Field(
        ..., description="Unique ID for the run"
    )
    agent_name: Optional[str] = Field(
        ..., description="Name of the agent"
    )
    task: Optional[str] = Field(
        ..., description="Task or query given to the agent"
    )
    output: Optional[str] = Field(
        ..., description="Output generated by the agent"
    )
    start_time: Optional[datetime] = Field(
        ..., description="Start time of the task"
    )
    end_time: Optional[datetime] = Field(
        ..., description="End time of the task"
    )
    duration: Optional[float] = Field(
        ...,
        description="Duration taken to complete the task (in seconds)",
    )


class MetadataSchema(BaseModel):
    swarm_id: Optional[str] = Field(
        generate_swarm_id(), description="Unique ID for the run"
    )
    task: Optional[str] = Field(
        ..., description="Task or query given to all agents"
    )
    description: Optional[str] = Field(
        "Concurrent execution of multiple agents",
        description="Description of the workflow",
    )
    agents: Optional[List[AgentOutputSchema]] = Field(
        ..., description="List of agent outputs and metadata"
    )
    timestamp: Optional[datetime] = Field(
        default_factory=datetime.now,
        description="Timestamp of the workflow execution",
    )


class ConcurrentWorkflow(BaseSwarm):
    """
    Represents a concurrent workflow that executes multiple agents concurrently in a production-grade manner.

    Args:
        name (str): The name of the workflow. Defaults to "ConcurrentWorkflow".
        description (str): The description of the workflow. Defaults to "Execution of multiple agents concurrently".
        agents (List[Agent]): The list of agents to be executed concurrently. Defaults to an empty list.
        metadata_output_path (str): The path to save the metadata output. Defaults to "agent_metadata.json".
        auto_save (bool): Flag indicating whether to automatically save the metadata. Defaults to False.
        output_schema (BaseModel): The output schema for the metadata. Defaults to MetadataSchema.
        max_loops (int): The maximum number of loops for each agent. Defaults to 1.
        return_str_on (bool): Flag indicating whether to return the output as a string. Defaults to False.
        agent_responses (list): The list of agent responses. Defaults to an empty list.
        auto_generate_prompts (bool): Flag indicating whether to auto-generate prompts for agents. Defaults to False.

    Raises:
        ValueError: If the list of agents is empty or if the description is empty.

    Attributes:
        name (str): The name of the workflow.
        description (str): The description of the workflow.
        agents (List[Agent]): The list of agents to be executed concurrently.
        metadata_output_path (str): The path to save the metadata output.
        auto_save (bool): Flag indicating whether to automatically save the metadata.
        output_schema (BaseModel): The output schema for the metadata.
        max_loops (int): The maximum number of loops for each agent.
        return_str_on (bool): Flag indicating whether to return the output as a string.
        agent_responses (list): The list of agent responses.
        auto_generate_prompts (bool): Flag indicating whether to auto-generate prompts for agents.
        retry_attempts (int): The number of retry attempts for failed agent executions.
        retry_wait_time (int): The initial wait time for retries in seconds.
    """

    def __init__(
        self,
        name: str = "ConcurrentWorkflow",
        description: str = "Execution of multiple agents concurrently",
        agents: List[Union[Agent, Callable]] = [],
        metadata_output_path: str = "agent_metadata.json",
        auto_save: bool = True,
        output_schema: BaseModel = MetadataSchema,
        max_loops: int = 1,
        return_str_on: bool = False,
        agent_responses: list = [],
        auto_generate_prompts: bool = False,
        output_type: OutputType = "dict",
        return_entire_history: bool = False,
        *args,
        **kwargs,
    ):
        super().__init__(
            name=name,
            description=description,
            agents=agents,
            *args,
            **kwargs,
        )
        self.name = name
        self.description = description
        self.agents = agents
        self.metadata_output_path = metadata_output_path
        self.auto_save = auto_save
        self.output_schema = output_schema
        self.max_loops = max_loops
        self.return_str_on = return_str_on
        self.agent_responses = agent_responses
        self.auto_generate_prompts = auto_generate_prompts
        self.max_workers = os.cpu_count()
        self.output_type = output_type
        self.return_entire_history = return_entire_history
        self.tasks = []  # Initialize tasks list

        self.reliability_check()

        self.conversation = Conversation()

    def disable_agent_prints(self):
        for agent in self.agents:
            agent.no_print = False

    def reliability_check(self):
        try:
            logger.info("Starting reliability checks")

            if self.name is None:
                logger.error("A name is required for the swarm")
                raise ValueError("A name is required for the swarm")

            if not self.agents:
                logger.error("The list of agents must not be empty.")
                raise ValueError(
                    "The list of agents must not be empty."
                )

            if not self.description:
                logger.error("A description is required.")
                raise ValueError("A description is required.")

            logger.info("Reliability checks completed successfully")
        except ValueError as e:
            logger.error(f"Reliability check failed: {e}")
            raise
        except Exception as e:
            logger.error(
                f"An unexpected error occurred during reliability checks: {e}"
            )
            raise

    def activate_auto_prompt_engineering(self):
        """
        Activates the auto-generate prompts feature for all agents in the workflow.

        Example:
            >>> workflow = ConcurrentWorkflow(agents=[Agent()])
            >>> workflow.activate_auto_prompt_engineering()
            >>> # All agents in the workflow will now auto-generate prompts.
        """
        if self.auto_generate_prompts is True:
            for agent in self.agents:
                agent.auto_generate_prompt = True

    # @retry(wait=wait_exponential(min=2), stop=stop_after_attempt(3))

    def transform_metadata_schema_to_str(
        self, schema: MetadataSchema
    ):
        """
        Converts the metadata swarm schema into a string format with the agent name, response, and time.

        Args:
            schema (MetadataSchema): The metadata schema to convert.

        Returns:
            str: The string representation of the metadata schema.

        Example:
            >>> metadata_schema = MetadataSchema()
            >>> metadata_str = workflow.transform_metadata_schema_to_str(metadata_schema)
            >>> print(metadata_str)
        """
        self.agent_responses = [
            f"Agent Name: {agent.agent_name}\nResponse: {agent.output}\n\n"
            for agent in schema.agents
        ]

        # Return the agent responses as a string
        return "\n".join(self.agent_responses)

    def save_metadata(self):
        """
        Saves the metadata to a JSON file based on the auto_save flag.

        Example:
            >>> workflow.save_metadata()
            >>> # Metadata will be saved to the specified path if auto_save is True.
        """
        # Save metadata to a JSON file
        if self.auto_save:
            logger.info(
                f"Saving metadata to {self.metadata_output_path}"
            )
            create_file_in_folder(
                os.getenv("WORKSPACE_DIR"),
                self.metadata_output_path,
                self.output_schema.model_dump_json(indent=4),
            )

    def _run(
        self, task: str, img: str = None, *args, **kwargs
    ) -> Union[Dict[str, Any], str]:
        """
        Runs the workflow for the given task, executes agents concurrently using ThreadPoolExecutor, and saves metadata.

        Args:
            task (str): The task or query to give to all agents.
            img (str): The image to be processed by the agents.

        Returns:
            Dict[str, Any]: The final metadata as a dictionary.
            str: The final metadata as a string if return_str_on is True.

        Example:
            >>> metadata = workflow.run(task="Example task", img="example.jpg")
            >>> print(metadata)
        """
        logger.info(
            f"Running concurrent workflow with {len(self.agents)} agents."
        )

        self.conversation.add(
            "User",
            task,
        )

        def run_agent(
            agent: Agent, task: str, img: str = None
        ) -> AgentOutputSchema:
            start_time = datetime.now()
            try:
                output = agent.run(task=task)

                self.conversation.add(
                    agent.agent_name,
                    output,
                )
            except Exception as e:
                logger.error(
                    f"Error running agent {agent.agent_name}: {e}"
                )
                raise

            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()

            agent_output = AgentOutputSchema(
                run_id=uuid.uuid4().hex,
                agent_name=agent.agent_name,
                task=task,
                output=output,
                start_time=start_time,
                end_time=end_time,
                duration=duration,
            )

            logger.info(
                f"Agent {agent.agent_name} completed task: {task} in {duration:.2f} seconds."
            )

            return agent_output

        with ThreadPoolExecutor(
            max_workers=os.cpu_count()
        ) as executor:
            agent_outputs = list(
                executor.map(
                    lambda agent: run_agent(agent, task), self.agents
                )
            )

        self.output_schema = MetadataSchema(
            swarm_id=uuid.uuid4().hex,
            task=task,
            description=self.description,
            agents=agent_outputs,
        )

        self.save_metadata()

        if self.return_str_on:
            return self.transform_metadata_schema_to_str(
                self.output_schema
            )

        elif self.return_entire_history:
            return self.conversation.return_history_as_string()
        elif self.output_type == "list":
            return self.conversation.return_messages_as_list()
        elif self.output_type == "dict":
            return self.conversation.return_messages_as_dictionary()
        else:
            return self.output_schema.model_dump_json(indent=4)

    def run(
        self,
        task: Optional[str] = None,
        img: Optional[str] = None,
        *args,
        **kwargs,
    ) -> Any:
        """
        Executes the agent's run method on a specified device.

        This method attempts to execute the agent's run method on a specified device, either CPU or GPU. It logs the device selection and the number of cores or GPU ID used. If the device is set to CPU, it can use all available cores or a specific core specified by `device_id`. If the device is set to GPU, it uses the GPU specified by `device_id`.

        Args:
            task (Optional[str], optional): The task to be executed. Defaults to None.
            img (Optional[str], optional): The image to be processed. Defaults to None.
            is_last (bool, optional): Indicates if this is the last task. Defaults to False.
            device (str, optional): The device to use for execution. Defaults to "cpu".
            device_id (int, optional): The ID of the GPU to use if device is set to "gpu". Defaults to 0.
            all_cores (bool, optional): If True, uses all available CPU cores. Defaults to True.
            all_gpus (bool, optional): If True, uses all available GPUS. Defaults to True.
            *args: Additional positional arguments to be passed to the execution method.
            **kwargs: Additional keyword arguments to be passed to the execution method.

        Returns:
            Any: The result of the execution.

        Raises:
            ValueError: If an invalid device is specified.
            Exception: If any other error occurs during execution.
        """
        if task is not None:
            self.tasks.append(task)

        try:
            outputs = self._run(task, img, *args, **kwargs)
            return outputs
        except ValueError as e:
            logger.error(f"Invalid device specified: {e}")
            raise e
        except Exception as e:
            logger.error(f"An error occurred during execution: {e}")
            raise e

    def run_batched(
        self, tasks: List[str]
    ) -> List[Union[Dict[str, Any], str]]:
        """
        Runs the workflow for a batch of tasks, executes agents concurrently for each task, and saves metadata in a production-grade manner.

        Args:
            tasks (List[str]): A list of tasks or queries to give to all agents.

        Returns:
            List[Union[Dict[str, Any], str]]: A list of final metadata for each task, either as a dictionary or a string.

        Example:
            >>> tasks = ["Task 1", "Task 2"]
            >>> results = workflow.run_batched(tasks)
            >>> print(results)
        """
        results = []
        for task in tasks:
            result = self.run(task)
            results.append(result)
        return results


# if __name__ == "__main__":
#     # Assuming you've already initialized some agents outside of this class
#     agents = [
#         Agent(
#             agent_name=f"Financial-Analysis-Agent-{i}",
#             system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
#             model_name="gpt-4o",
#             max_loops=1,
#         )
#         for i in range(3)  # Adjust number of agents as needed
#     ]

#     # Initialize the workflow with the list of agents
#     workflow = ConcurrentWorkflow(
#         agents=agents,
#         metadata_output_path="agent_metadata_4.json",
#         return_str_on=True,
#     )

#     # Define the task for all agents
#     task = "How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria?"

#     # Run the workflow and save metadata
#     metadata = workflow.run(task)
#     print(metadata)
