from autogen_core import Component, ComponentModel
from pydantic import BaseModel
from typing_extensions import Self

from autogen_agentchat.agents import BaseChatAgent

from ._task_runner_tool import TaskRunnerTool


class AgentToolConfig(BaseModel):
    """Configuration for the AgentTool."""

    agent: ComponentModel
    """The agent to be used for running the task."""

    return_value_as_last_message: bool = False
    """Whether to return the value as the last message of the task result."""


class AgentTool(TaskRunnerTool, Component[AgentToolConfig]):
    """Tool that can be used to run a task using an agent.

    The tool returns the result of the task execution as a :class:`~autogen_agentchat.base.TaskResult` object.

    .. important::
        When using AgentTool, you **must** disable parallel tool calls in the model client configuration
        to avoid concurrency issues. Agents cannot run concurrently as they maintain internal state
        that would conflict with parallel execution. For example, set ``parallel_tool_calls=False``
        for :class:`~autogen_ext.models.openai.OpenAIChatCompletionClient` and
        :class:`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`.

    Args:
        agent (BaseChatAgent): The agent to be used for running the task.
        return_value_as_last_message (bool): Whether to use the last message content of the task result
            as the return value of the tool in :meth:`~autogen_agentchat.tools.TaskRunnerTool.return_value_as_string`.
            If set to True, the last message content will be returned as a string.
            If set to False, the tool will return all messages in the task result as a string concatenated together,
            with each message prefixed by its source (e.g., "writer: ...", "assistant: ...").

    Example:

        .. code-block:: python

            import asyncio

            from autogen_agentchat.agents import AssistantAgent
            from autogen_agentchat.tools import AgentTool
            from autogen_agentchat.ui import Console
            from autogen_ext.models.openai import OpenAIChatCompletionClient


            async def main() -> None:
                model_client = OpenAIChatCompletionClient(model="gpt-4.1")
                writer = AssistantAgent(
                    name="writer",
                    description="A writer agent for generating text.",
                    model_client=model_client,
                    system_message="Write well.",
                )
                writer_tool = AgentTool(agent=writer)

                # Create model client with parallel tool calls disabled for the main agent
                main_model_client = OpenAIChatCompletionClient(model="gpt-4.1", parallel_tool_calls=False)
                assistant = AssistantAgent(
                    name="assistant",
                    model_client=main_model_client,
                    tools=[writer_tool],
                    system_message="You are a helpful assistant.",
                )
                await Console(assistant.run_stream(task="Write a poem about the sea."))


            asyncio.run(main())
    """

    component_config_schema = AgentToolConfig
    component_provider_override = "autogen_agentchat.tools.AgentTool"

    def __init__(self, agent: BaseChatAgent, return_value_as_last_message: bool = False) -> None:
        self._agent = agent
        super().__init__(
            agent, agent.name, agent.description, return_value_as_last_message=return_value_as_last_message
        )

    def _to_config(self) -> AgentToolConfig:
        return AgentToolConfig(
            agent=self._agent.dump_component(),
            return_value_as_last_message=self._return_value_as_last_message,
        )

    @classmethod
    def _from_config(cls, config: AgentToolConfig) -> Self:
        return cls(BaseChatAgent.load_component(config.agent), config.return_value_as_last_message)
