import concurrent.futures
import json
import traceback
from typing import Callable, List, Optional

from loguru import logger
from pydantic import BaseModel, Field

from swarms.structs.conversation import Conversation
from swarms.structs.ma_blocks import find_agent_by_name
from swarms.tools.base_tool import BaseTool
from swarms.utils.formatter import formatter
from swarms.utils.history_output_formatter import (
    history_output_formatter,
)
from swarms.utils.litellm_wrapper import LiteLLM
from swarms.utils.output_types import OutputType
from swarms.telemetry.otel import (
    ContextThreadPoolExecutor,
    capture_init,
    trace_run,
)
from swarms.utils.generate_id import generate_id


class HandOffsResponse(BaseModel):
    """
    Response from the boss agent indicating which agent should handle the task.

    This model encapsulates the reasoning behind the agent selection, the name of the selected agent, and an optional modified version of the task. It is used to communicate the boss agent's decision and rationale to the rest of the system.
    """

    reasoning: str = Field(
        description="A detailed explanation for why this agent was selected. This should include the logic or criteria used by the boss agent to make the selection, providing transparency and traceability for the routing decision."
    )
    agent_name: str = Field(
        description="The name of the agent selected to handle the task. This should match the identifier of one of the available agents in the system, ensuring the task is routed correctly."
    )
    task: Optional[str] = Field(
        None,
        description="An optional, modified version of the original task. If the boss agent determines that the task should be rephrased or clarified before execution, the new version is provided here. If no modification is needed, this field may be left as None.",
    )


class MultipleHandOffsResponse(BaseModel):
    """
    Response from the boss agent indicating which agents should handle the task.
    """

    handoffs: List[HandOffsResponse] = Field(
        description="A list of handoffs, each containing the reasoning, agent name, and task for each agent."
    )


def get_agent_response_schema(model_name: str = None):
    """
    Return the JSON schema for the boss agent's structured handoff response.

    Args:
        model_name (str, optional): Reserved for future model-specific schema
            tweaks. Currently unused.

    Returns:
        dict: A JSON schema dict derived from ``MultipleHandOffsResponse`` that
        can be passed to a structured-output model as the response format.
    """
    return BaseTool().base_model_to_dict(MultipleHandOffsResponse)


def agent_boss_router_prompt(agent_descriptions: any):
    """
    Build the system prompt that instructs the boss agent how to route tasks.

    The prompt embeds the rendered ``agent_descriptions`` block so the boss
    model knows what agents are available and how to choose between them.

    Args:
        agent_descriptions: A string (or string-coercible value) listing the
            available agents, typically in the form
            ``"- <agent_name>: <description>"`` joined by newlines.

    Returns:
        str: The fully rendered boss-agent system prompt.
    """
    return f"""
        You are an intelligent boss agent responsible for routing tasks to the most appropriate specialized agents.

        Available agents:
        {agent_descriptions}

        Your primary responsibilities:
        1. **Understand the user's intent and task requirements** - Carefully analyze what the user is asking for
        2. **Determine task complexity** - Identify if this is a single task or multiple related tasks
        3. **Match capabilities to requirements** - Select the most suitable agent(s) based on their descriptions and expertise
        4. **Provide clear reasoning** - Explain your selection logic transparently
        5. **Optimize task assignment** - Modify tasks if needed to better suit the selected agent's capabilities

        **Routing Logic:**
        - **Single Task**: If the user presents one clear task, select the single best agent for that task
        - **Multiple Tasks**: If the user presents multiple distinct tasks or a complex task that requires different expertise areas, select multiple agents and break down the work accordingly
        - **Collaborative Tasks**: If a task benefits from multiple perspectives or requires different skill sets, assign it to multiple agents

        **Response Format:**
        You must respond with JSON containing a "handoffs" array with one or more handoff objects. Each handoff object must contain:
        - reasoning: Detailed explanation of why this agent was selected and how they fit the task requirements
        - agent_name: Name of the chosen agent (must exactly match one of the available agents)
        - task: A modified/optimized version of the task if needed, or the original task if no modification is required

        **Important Guidelines:**
        - Always analyze the user's intent first before making routing decisions
        - Consider task dependencies and whether agents need to work sequentially or in parallel
        - If multiple agents are selected, ensure their tasks are complementary and don't conflict
        - Provide specific, actionable reasoning for each agent selection
        - Ensure all agent names exactly match the available agents list
        - If a single agent can handle the entire request efficiently, use only one agent
        - If the request requires different expertise areas, use multiple agents with clearly defined, non-overlapping tasks

        Remember: Your goal is to maximize task completion quality by matching the right agent(s) to the right work based on their capabilities and the user's actual needs.
        """


class MultiAgentRouter:
    """
    Routes tasks to appropriate agents based on their capabilities.

    This class is responsible for managing a pool of agents and routing incoming tasks to the most suitable agent. It uses a boss agent to analyze the task and select the best agent for the job. The boss agent's decision is based on the capabilities and descriptions of the available agents.

    Attributes:
        name (str): The name of the router.
        description (str): A description of the router's purpose.
        agents (List[Callable]): List of agents managed by the router. Name lookups go through `find_agent_by_name`.
        model (str): The model to use for the boss agent.
        temperature (float): The temperature for the boss agent's model.
        shared_memory_system (callable): A shared memory system for agents to query.
        output_type (OutputType): The type of output expected from the agents.
        print_on (bool): Whether to print the boss agent's decision.
        system_prompt (str): Custom system prompt for the router.
        skip_null_tasks (bool): Whether to skip executing agents when their assigned task is null or None.
        conversation (Conversation): The conversation history.
        function_caller (LiteLLM): An instance of LiteLLM for calling the boss agent.
    """

    def __init__(
        self,
        id: Optional[str] = None,
        name: str = "swarm-router",
        description: str = "Routes tasks to specialized agents based on their capabilities",
        agents: List[Callable] = None,
        model: str = "gpt-5.4",
        temperature: float = 0.1,
        shared_memory_system: callable = None,
        output_type: OutputType = "dict",
        print_on: bool = True,
        system_prompt: str = None,
        skip_null_tasks: bool = True,
        *args,
        **kwargs,
    ):
        """
        Initializes the MultiAgentRouter with a list of agents and configuration options.

        Args:
            name (str, optional): The name of the router. Defaults to "swarm-router".
            description (str, optional): A description of the router's purpose. Defaults to "Routes tasks to specialized agents based on their capabilities".
            agents (List[Agent], optional): A list of agents to be managed by the router. Defaults to an empty list.
            model (str, optional): The model to use for the boss agent. Defaults to "gpt-5.4".
            temperature (float, optional): The temperature for the boss agent's model. Defaults to 0.1.
            shared_memory_system (callable, optional): A shared memory system for agents to query. Defaults to None.
            output_type (OutputType, optional): The type of output expected from the agents. Defaults to "dict".
            print_on (bool, optional): Whether to print the boss agent's decision. Defaults to True.
            system_prompt (str, optional): Custom system prompt for the router. Defaults to None.
            skip_null_tasks (bool, optional): Whether to skip executing agents when their assigned task is null or None. Defaults to True.
        """
        self.id = id or generate_id("multi-agent-router")
        self.name = name
        self.description = description
        self.shared_memory_system = shared_memory_system
        self.output_type = output_type
        self.model = model
        self.temperature = temperature
        self.print_on = print_on
        self.system_prompt = system_prompt
        self.skip_null_tasks = skip_null_tasks
        self.agents = list(agents) if agents else []
        self.conversation = Conversation()

        router_system_prompt = ""

        router_system_prompt += (
            self.system_prompt if self.system_prompt else ""
        )
        router_system_prompt += self._create_boss_system_prompt()

        self.function_caller = LiteLLM(
            model_name=self.model,
            system_prompt=router_system_prompt,
            temperature=self.temperature,
            tool_choice="auto",
            parallel_tool_calls=True,
            response_format=MultipleHandOffsResponse,
            *args,
            **kwargs,
        )

        # Capture the full __init__ configuration if telemetry is enabled.
        capture_init(self)

    def __repr__(self):
        """
        Return a debug-friendly representation listing the router name and the
        names of every agent it manages.

        Returns:
            str: A repr of the form
            ``MultiAgentRouter(name=<name>, agents=[<names>])``.
        """
        return f"MultiAgentRouter(name={self.name}, agents={[a.agent_name for a in self.agents]})"

    def query_ragent(self, task: str) -> str:
        """
        Query the shared memory / research-agent system attached to the router.

        Args:
            task (str): The query string to forward to ``shared_memory_system``.

        Returns:
            str: The response returned by the shared memory system.

        Raises:
            AttributeError: If no ``shared_memory_system`` was configured on the
                router.
        """
        return self.shared_memory_system.query(task)

    def _create_boss_system_prompt(self) -> str:
        """
        Creates a system prompt for the boss agent that includes information about all available agents.

        Returns:
            str: The system prompt for the boss agent.
        """
        agent_descriptions = "\n".join(
            [
                f"- {agent.agent_name}: {agent.description}"
                for agent in self.agents
            ]
        )

        return agent_boss_router_prompt(agent_descriptions)

    def handle_single_handoff(
        self, boss_response_str: dict, task: str
    ) -> dict:
        """
        Execute the single agent selected by the boss and record its response.

        Looks up the agent named in the first (and only) handoff, runs it on
        the modified task (falling back to the original task if the boss did
        not rewrite it), and appends the agent's response to
        ``self.conversation``. When ``skip_null_tasks`` is True and the
        resolved task is empty or ``None``, execution is skipped.

        Args:
            boss_response_str (dict): Parsed boss decision. Expected to contain
                a ``"handoffs"`` list with at least one entry shaped like
                ``{"agent_name": str, "task": Optional[str], ...}``.
            task (str): The original user task, used as a fallback when the
                boss did not supply a rewritten task.

        Returns:
            dict: Currently returns ``None`` implicitly; the conversation is
            mutated in place. Kept as ``dict`` in the signature for parity
            with related methods.

        Raises:
            ValueError: If the boss selected an agent name that is not
                registered with this router.
        """

        handoff_agent_name = boss_response_str["handoffs"][0][
            "agent_name"
        ]

        selected_agent = find_agent_by_name(
            self.agents, handoff_agent_name
        )

        # Use the modified task if provided, otherwise use original task
        final_task = boss_response_str["handoffs"][0]["task"] or task

        # Skip execution if task is null/None and skip_null_tasks is True
        if self.skip_null_tasks and (
            final_task is None or final_task == ""
        ):
            if self.print_on:
                logger.info(
                    f"Skipping execution for agent {selected_agent.agent_name} - task is null/None"
                )

        # Use the agent's run method directly
        agent_response = selected_agent.run(final_task)

        self.conversation.add(
            role=selected_agent.agent_name, content=agent_response
        )

        # return agent_response

    def handle_multiple_handoffs(
        self, boss_response_str: dict, task: str
    ) -> dict:
        """
        Execute every agent selected by the boss and record the first response.

        Validates that every ``agent_name`` referenced in ``boss_response_str``
        exists before running anything. Each agent is run with its boss-supplied
        task (or the original ``task`` when the boss did not rewrite it). When
        ``skip_null_tasks`` is True, agents whose resolved task is empty or
        ``None`` are skipped. After execution, the first selected agent's
        response is appended to ``self.conversation``.

        Args:
            boss_response_str (dict): Parsed boss decision containing a
                ``"handoffs"`` list, each entry shaped like
                ``{"agent_name": str, "task": Optional[str], ...}``.
            task (str): The original user task, used as a fallback when the
                boss did not supply a rewritten task.

        Returns:
            dict: Currently returns ``None`` implicitly; the conversation is
            mutated in place. Kept as ``dict`` in the signature for parity
            with related methods.

        Raises:
            ValueError: If any of the boss-selected agents is not registered
                with this router.
        """

        # Validate that the selected agents exist
        for handoff in boss_response_str["handoffs"]:
            try:
                find_agent_by_name(self.agents, handoff["agent_name"])
            except (TypeError, ValueError):
                raise ValueError(
                    f"Boss selected unknown agent: {handoff['agent_name']}"
                )

        # Get the selected agents and their tasks
        selected_agents = []
        final_tasks = []
        skipped_agents = []

        for handoff in boss_response_str["handoffs"]:
            agent = find_agent_by_name(
                self.agents, handoff["agent_name"]
            )
            final_task = handoff["task"] or task

            # Skip execution if task is null/None and skip_null_tasks is True
            if self.skip_null_tasks and (
                final_task is None or final_task == ""
            ):
                if self.print_on:
                    logger.info(
                        f"Skipping execution for agent {agent.agent_name} - task is null/None"
                    )
                skipped_agents.append(agent.agent_name)
                continue

            selected_agents.append(agent)
            final_tasks.append(final_task)

        # Execute agents only if there are valid tasks
        if selected_agents:
            # Use the agents' run method directly
            agent_responses = [
                agent.run(final_task)
                for agent, final_task in zip(
                    selected_agents, final_tasks
                )
            ]

            self.conversation.add(
                role=selected_agents[0].agent_name,
                content=agent_responses[0],
            )

        # return agent_responses

    def route_task(self, task: str) -> dict:
        """
        Routes a task to the appropriate agent and returns their response.

        Args:
            task (str): The task to be routed.

        Returns:
            dict: A dictionary containing the routing result, including the selected agent, reasoning, and response.
        """
        try:
            self.conversation.add(role="user", content=task)

            # Get boss decision using function calling
            boss_response_str = self.function_caller.run(task)

            boss_response_str = json.loads(boss_response_str)

            if self.print_on:
                formatter.print_panel(
                    json.dumps(boss_response_str, indent=4),
                    title=self.name,
                )

            if len(boss_response_str["handoffs"]) > 1:
                self.handle_multiple_handoffs(boss_response_str, task)
            else:
                self.handle_single_handoff(boss_response_str, task)

            return history_output_formatter(
                conversation=self.conversation, type=self.output_type
            )

        except Exception as e:
            logger.error(
                f"Error routing task: {str(e)} Traceback: {traceback.format_exc()}"
            )
            raise

    @trace_run(
        "MultiAgentRouter.run",
        input_params=("task", "tasks", "img", "imgs"),
    )
    def run(self, task: str):
        """
        Route a single task through the boss agent and execute the selected
        worker(s).

        Args:
            task (str): The task to route.

        Returns:
            Any: The formatted conversation history, shaped according to
            ``self.output_type`` (see ``history_output_formatter``).
        """
        return self.route_task(task)

    def __call__(self, task: str):
        """
        Convenience alias for :meth:`run` so the router can be invoked
        directly: ``router(task)``.

        Args:
            task (str): The task to route.

        Returns:
            Any: The formatted conversation history, shaped according to
            ``self.output_type``.
        """
        return self.route_task(task)

    def batch_run(self, tasks: List[str] = []):
        """
        Route a batch of tasks sequentially.

        Each task is independently routed; failures are logged and the
        corresponding result is omitted from the returned list (the batch
        does not abort on a single failure).

        Args:
            tasks (List[str]): Tasks to route in order.

        Returns:
            List[Any]: Routed results for the tasks that succeeded, in input
            order.
        """
        results = []
        for task in tasks:
            try:
                result = self.route_task(task)
                results.append(result)
            except Exception as e:
                logger.error(f"Error routing task: {str(e)}")
        return results

    def concurrent_batch_run(self, tasks: List[str] = []):
        """
        Route a batch of tasks in parallel using a thread pool.

        Tasks are dispatched to a ``ThreadPoolExecutor`` and results are
        collected as they complete. Order of returned results is **not
        guaranteed** to match input order. Failures are logged and silently
        dropped from the result list.

        Args:
            tasks (List[str]): Tasks to route concurrently.

        Returns:
            List[Any]: Routed results for the tasks that succeeded, in
            completion order.
        """
        results = []
        with ContextThreadPoolExecutor() as executor:
            futures = [
                executor.submit(self.route_task, task)
                for task in tasks
            ]
            for future in concurrent.futures.as_completed(futures):
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    logger.error(f"Error routing task: {str(e)}")
        return results


# # Example usage:
# if __name__ == "__main__":
#     # Define some example agents
#     agents = [
#         Agent(
#             agent_name="ResearchAgent",
#             description="Specializes in researching topics and providing detailed, factual information",
#             system_prompt="You are a research specialist. Provide detailed, well-researched information about any topic, citing sources when possible.",
#             model_name="openai/gpt-4o",
#         ),
#         Agent(
#             agent_name="CodeExpertAgent",
#             description="Expert in writing, reviewing, and explaining code across multiple programming languages",
#             system_prompt="You are a coding expert. Write, review, and explain code with a focus on best practices and clean code principles.",
#             model_name="openai/gpt-4o",
#         ),
#         Agent(
#             agent_name="WritingAgent",
#             description="Skilled in creative and technical writing, content creation, and editing",
#             system_prompt="You are a writing specialist. Create, edit, and improve written content while maintaining appropriate tone and style.",
#             model_name="openai/gpt-4o",
#         ),
#     ]

#     # Initialize router
#     router = MultiAgentRouter(agents=agents)

#     # Example task
#     task = "Write a Python function to calculate fibonacci numbers"

#     try:
#         # Process the task
#         result = router.route_task(task)
#         print(f"Selected Agent: {result['boss_decision']['selected_agent']}")
#         print(f"Reasoning: {result['boss_decision']['reasoning']}")
#         print(f"Total Time: {result['total_time']:.2f}s")

#     except Exception as e:
#         print(f"Error occurred: {str(e)}")
