import asyncio
import concurrent.futures
import os
from typing import Any, Callable, List, Optional, Union

from loguru import logger

from swarms.structs.agent import Agent
from swarms.structs.omni_agent_types import AgentType
from swarms.telemetry.otel import ContextThreadPoolExecutor
from swarms.utils.get_cpu_cores import max_workers_95_percent


def run_single_agent(
    agent: AgentType, task: str, *args, **kwargs
) -> Any:
    """
    Run a single agent synchronously with the given task.

    This function provides a synchronous wrapper for executing a single agent
    with a specific task. It passes through any additional arguments and
    keyword arguments to the agent's run method.

    Args:
        agent (AgentType): The agent instance to execute
        task (str): The task string to be executed by the agent
        *args: Variable length argument list passed to agent.run()
        **kwargs: Arbitrary keyword arguments passed to agent.run()

    Returns:
        Any: The result returned by the agent's run method

    Example:
        >>> agent = SomeAgent()
        >>> result = run_single_agent(agent, "Analyze this data")
        >>> print(result)
    """
    return agent.run(task=task, *args, **kwargs)


async def run_agent_async(agent: AgentType, task: str) -> Any:
    """
    Run an agent asynchronously using asyncio event loop.

    This function executes a single agent asynchronously by running it in a
    thread executor to avoid blocking the event loop. It's designed to be
    used within async contexts for concurrent execution.

    Args:
        agent (AgentType): The agent instance to execute asynchronously
        task (str): The task string to be executed by the agent

    Returns:
        Any: The result returned by the agent's run method

    Example:
        >>> async def main():
        ...     agent = SomeAgent()
        ...     result = await run_agent_async(agent, "Process data")
        ...     return result
    """
    return await asyncio.to_thread(run_single_agent, agent, task)


async def run_agents_concurrently_async(
    agents: List[AgentType], task: str
) -> List[Any]:
    """
    Run multiple agents concurrently using asyncio gather.

    This function executes multiple agents concurrently using asyncio.gather(),
    which runs all agents in parallel and waits for all to complete. Each agent
    runs the same task asynchronously.

    Args:
        agents (List[AgentType]): List of agent instances to run concurrently
        task (str): The task string to be executed by all agents

    Returns:
        List[Any]: List of results from each agent in the same order as input

    Example:
        >>> async def main():
        ...     agents = [Agent1(), Agent2(), Agent3()]
        ...     results = await run_agents_concurrently_async(agents, "Analyze data")
        ...     for i, result in enumerate(results):
        ...         print(f"Agent {i+1} result: {result}")
    """
    results = await asyncio.gather(
        *(run_agent_async(agent, task) for agent in agents)
    )
    return results


def run_agents_concurrently(
    agents: List["AgentType"],
    task: str,
    img: Optional[str] = None,
    max_workers: Optional[int] = None,
    return_agent_output_dict: bool = False,
) -> Any:
    """
    Execute multiple agents concurrently using a ThreadPoolExecutor.

    This function runs agent tasks in parallel threads, benefitting I/O-bound or mixed-load scenarios.
    Each agent receives the same 'task' (and optional 'img' argument) and runs its .run() method.
    The number of worker threads defaults to 95% of the available CPU cores, unless otherwise specified.

    Args:
        agents (List[AgentType]): List of agent instances to execute concurrently.
        task (str): Task string to pass to all agent run() methods.
        img (Optional[str]): Optional image data to pass to agent run() if supported.
        max_workers (Optional[int]): Maximum threads for the executor (default: 95% of CPU cores).
        return_agent_output_dict (bool): If True, returns a dict mapping agent names to outputs.
                                         Otherwise returns a list of results in completion order.

    Returns:
        List[Any] or Dict[str, Any]: List of results from each agent's run() method in completion order,
                                     or a dict of agent names to results (preserving agent order)
                                     if return_agent_output_dict is True.
                                     If an agent fails, the corresponding result is the Exception.

    Notes:
        - ThreadPoolExecutor is used for efficient, parallel execution.
        - By default, utilizes nearly all available CPU cores for optimal performance.
        - Any Exception during agent execution is caught and included in the results.
        - If return_agent_output_dict is True, the results dict preserves agent input order.
        - Otherwise, the results list is in order of completion (not input order).

    Example:
        >>> agents = [Agent1(), Agent2()]
        >>> # As list
        >>> results = run_agents_concurrently(agents, task="Summarize", img=None)
        >>> # As dict
        >>> results_dict = run_agents_concurrently(
        ...    agents, task="Summarize", return_agent_output_dict=True)
        >>> for name, val in results_dict.items():
        ...     print(f"Result from {name}: {val}")
    """
    try:
        if max_workers is None:
            max_workers = max_workers_95_percent()

        futures = []
        agent_id_map = {}

        with ContextThreadPoolExecutor(
            max_workers=max_workers
        ) as executor:
            for agent in agents:
                agent_kwargs = {}
                if task is not None:
                    agent_kwargs["task"] = task
                if img is not None:
                    agent_kwargs["img"] = img
                future = executor.submit(agent.run, **agent_kwargs)
                futures.append(future)
                agent_id_map[future] = agent

            if return_agent_output_dict:
                # Use agent name as key, preserve input order
                output_dict = {}
                for agent, future in zip(agents, futures):
                    try:
                        result = future.result()
                    except Exception as e:
                        result = e
                    # Prefer .agent_name or .name, fallback to str(agent)
                    name = (
                        getattr(agent, "agent_name", None)
                        or getattr(agent, "name", None)
                        or str(agent)
                    )
                    output_dict[name] = result
                return output_dict
            else:
                results = []
                for future in concurrent.futures.as_completed(
                    futures
                ):
                    try:
                        result = future.result()
                        results.append(result)
                    except Exception as e:
                        results.append(e)
                return results

    except Exception as e:
        logger.error(
            f"Error running_agents_concurrently: {e} Traceback: {e.__traceback__}"
        )
        raise e


def run_agents_concurrently_multiprocess(
    agents: List[Agent], task: str, batch_size: int = os.cpu_count()
) -> List[Any]:
    """
    Run multiple agents concurrently in batches using asyncio for optimized performance.

    This function processes agents in batches to avoid overwhelming system resources
    while still achieving high concurrency. It uses asyncio internally to manage
    the concurrent execution of agent batches.

    Args:
        agents (List[Agent]): List of Agent instances to run concurrently
        task (str): The task string to be executed by all agents
        batch_size (int, optional): Number of agents to run in parallel in each batch.
                                   Defaults to the number of CPU cores for optimal resource usage

    Returns:
        List[Any]: List of results from each agent, maintaining the order of input agents

    Note:
        - Processes agents in batches to prevent resource exhaustion
        - Uses asyncio for efficient concurrent execution within batches
        - Results are returned in the same order as input agents

    Example:
        >>> agents = [Agent1(), Agent2(), Agent3(), Agent4(), Agent5()]
        >>> results = run_agents_concurrently_multiprocess(agents, "Analyze data", batch_size=2)
        >>> print(f"Processed {len(results)} agents")
    """
    results = []

    # Use a dedicated event loop to ensure compatibility across contexts.
    loop = asyncio.new_event_loop()
    try:
        # Process agents in batches to avoid overwhelming system resources
        for i in range(0, len(agents), batch_size):
            batch = agents[i : i + batch_size]
            batch_results = loop.run_until_complete(
                run_agents_concurrently_async(batch, task)
            )
            results.extend(batch_results)
    finally:
        loop.close()

    return results


def batched_grid_agent_execution(
    agents: List["AgentType"],
    tasks: List[str],
    max_workers: int = None,
) -> List[Any]:
    """
    Run multiple agents with different tasks concurrently using ThreadPoolExecutor.

    This function pairs each agent with a specific task and executes them concurrently.
    It's designed for scenarios where different agents need to work on different tasks
    simultaneously, creating a grid-like execution pattern.

    Args:
        agents (List[AgentType]): List of agent instances to execute
        tasks (List[str]): List of task strings, one for each agent. Must match the number of agents
        max_workers (int, optional): Maximum number of threads to use.
                                   Defaults to 90% of available CPU cores for optimal performance

    Returns:
        List[Any]: List of results from each agent in the same order as input agents.
                  If an agent fails, the exception is included in the results.

    Raises:
        ValueError: If the number of agents doesn't match the number of tasks

    Note:
        - Uses 90% of CPU cores by default for optimal resource utilization
        - Results maintain the same order as input agents
        - Handles exceptions gracefully by including them in results

    Example:
        >>> agents = [Agent1(), Agent2(), Agent3()]
        >>> tasks = ["Task A", "Task B", "Task C"]
        >>> results = batched_grid_agent_execution(agents, tasks)
        >>> for i, result in enumerate(results):
        ...     print(f"Agent {i+1} with {tasks[i]}: {result}")
    """

    if max_workers is None:
        max_workers = max_workers_95_percent()

    logger.info(
        f"Batch Grid Execution with {len(agents)} agents and number of tasks: {len(tasks)}"
    )

    if len(agents) != len(tasks):
        raise ValueError(
            "The number of agents must match the number of tasks."
        )

    results = [None] * len(agents)
    with ContextThreadPoolExecutor(
        max_workers=max_workers
    ) as executor:
        future_to_index = {
            executor.submit(run_single_agent, agent, task): idx
            for idx, (agent, task) in enumerate(zip(agents, tasks))
        }
        for future in concurrent.futures.as_completed(
            future_to_index
        ):
            idx = future_to_index[future]
            try:
                results[idx] = future.result()
            except Exception as e:
                results[idx] = e

    return results


def run_agents_with_different_tasks(
    agent_task_pairs: List[tuple["AgentType", str]],
    batch_size: int = 10,
    max_workers: int = None,
) -> List[Any]:
    """
    Run multiple agents with different tasks concurrently, processing them in batches.

    This function executes each agent on its corresponding task, processing the agent-task pairs
    in batches for efficient resource utilization. It's designed for scenarios where you have
    a large number of agent-task pairs that need to be processed efficiently.

    Args:
        agent_task_pairs (List[tuple[AgentType, str]]): List of (agent, task) tuples to execute.
                                                        Each tuple contains an agent instance and its task
        batch_size (int, optional): Number of agent-task pairs to process in parallel in each batch.
                                   Defaults to 10 for balanced resource usage
        max_workers (int, optional): Maximum number of threads to use for each batch.
                                   If None, uses the default from batched_grid_agent_execution

    Returns:
        List[Any]: List of outputs from each agent-task pair, maintaining the same order as input pairs.
                  If an agent fails, the exception is included in the results.

    Note:
        - Processes agent-task pairs in batches to prevent resource exhaustion
        - Results maintain the same order as input pairs
        - Handles exceptions gracefully by including them in results
        - Uses batched_grid_agent_execution internally for each batch

    Example:
        >>> pairs = [(agent1, "Task A"), (agent2, "Task B"), (agent3, "Task C")]
        >>> results = run_agents_with_different_tasks(pairs, batch_size=5)
        >>> for i, result in enumerate(results):
        ...     agent, task = pairs[i]
        ...     print(f"Agent {agent.agent_name} with {task}: {result}")
    """
    if max_workers is None:
        max_workers = max_workers_95_percent()

    if not agent_task_pairs:
        return []

    results = []
    total_pairs = len(agent_task_pairs)
    for i in range(0, total_pairs, batch_size):
        batch = agent_task_pairs[i : i + batch_size]
        agents, tasks = zip(*batch)
        batch_results = batched_grid_agent_execution(
            list(agents), list(tasks), max_workers=max_workers
        )
        results.extend(batch_results)
    return results


def get_swarms_info(swarms: List[Callable]) -> str:
    """
    Fetch and format information about all available swarms in the system.

    This function provides a comprehensive overview of all swarms currently
    available in the system, including their names, descriptions, agent counts,
    and swarm types. It's useful for debugging, monitoring, and system introspection.

    Args:
        swarms (List[Callable]): List of swarm instances to get information about.
                               Each swarm should have name, description, agents, and swarm_type attributes

    Returns:
        str: A formatted string containing detailed information about all swarms.
             Returns "No swarms currently available in the system." if the list is empty.

    Note:
        - Each swarm is expected to have the following attributes:
          - name: The name of the swarm
          - description: A description of the swarm's purpose
          - agents: A list of agents in the swarm
          - swarm_type: The type/category of the swarm
        - The output is formatted for human readability with clear section headers

    Example:
        >>> swarms = [swarm1, swarm2, swarm3]
        >>> info = get_swarms_info(swarms)
        >>> print(info)
        Available Swarms:

        [Swarm 1]
        Name: Data Processing Swarm
        Description: Handles data analysis tasks
        Length of Agents: 5
        Swarm Type: Analysis
        ...
    """
    if not swarms:
        return "No swarms currently available in the system."

    swarm_info = [
        "Available Swarms:",
        "",
    ]  # Empty string for line spacing

    for idx, swarm in enumerate(swarms, 1):
        swarm_info.extend(
            [
                f"[Swarm {idx}]",
                f"Name: {swarm.name}",
                f"Description: {swarm.description}",
                f"Length of Agents: {len(swarm.agents)}",
                f"Swarm Type: {swarm.swarm_type}",
                "",  # Empty string for line spacing between swarms
            ]
        )

    return "\n".join(swarm_info).strip()


def get_agents_info(
    agents: List[Union[Agent, Callable]], team_name: str = None
) -> str:
    """
    Fetch and format information about all available agents in the system.

    This function provides a comprehensive overview of all agents currently
    available in the system, including their names, descriptions, roles,
    models, and configuration details. It's useful for debugging, monitoring,
    and system introspection.

    Args:
        agents (List[Union[Agent, Callable]]): List of agent instances to get information about.
                                             Each agent should have agent_name, agent_description,
                                             role, model_name, and max_loops attributes
        team_name (str, optional): Optional team name to include in the output header.
                                 If None, uses a generic header

    Returns:
        str: A formatted string containing detailed information about all agents.
             Returns "No agents currently available in the system." if the list is empty.

    Note:
        - Each agent is expected to have the following attributes:
          - agent_name: The name of the agent
          - agent_description: A description of the agent's purpose
          - role: The role or function of the agent
          - model_name: The AI model used by the agent
          - max_loops: The maximum number of loops the agent can execute
        - The output is formatted for human readability with clear section headers
        - Team name is included in the header if provided

    Example:
        >>> agents = [agent1, agent2, agent3]
        >>> info = get_agents_info(agents, team_name="Data Team")
        >>> print(info)
        Available Agents for Team: Data Team

        [Agent 1]
        Name: Data Analyzer
        Description: Analyzes data patterns
        Role: Analyst
        Model: gpt-4
        Max Loops: 10
        ...
    """
    if not agents:
        return "No agents currently available in the system."

    agent_info = [
        f"Available Agents for Team: {team_name}",
        "",
    ]  # Empty string for line spacing

    for idx, agent in enumerate(agents, 1):
        agent_info.extend(
            [
                "\n",
                f"[Agent {idx}]",
                f"Name: {agent.agent_name}",
                f"Description: {agent.agent_description}",
                f"Role: {agent.role}",
                f"Model: {agent.model_name}",
                f"Max Loops: {agent.max_loops}",
                "\n",
            ]
        )

    return "\n".join(agent_info).strip()
