# -*- coding: utf-8 -*-
"""The toolkit class for tool calls in agentscope.

TODO: We should consider to split this `Toolkit` class in the future.
"""
# pylint: disable=too-many-lines
import asyncio
import inspect
import os
from copy import deepcopy
from functools import partial, wraps
from typing import (
    AsyncGenerator,
    Literal,
    Any,
    Type,
    Generator,
    Callable,
    Awaitable,
    Coroutine,
)

import mcp
import shortuuid
from pydantic import (
    BaseModel,
    Field,
    create_model,
)

from ._async_wrapper import (
    _async_generator_wrapper,
    _object_wrapper,
    _sync_generator_wrapper,
)
from ._response import ToolResponse
from ._types import ToolGroup, AgentSkill, RegisteredToolFunction
from .._utils._common import _parse_tool_function
from ..mcp import (
    MCPToolFunction,
    MCPClientBase,
    StatefulClientBase,
)
from ..message import (
    ToolUseBlock,
    TextBlock,
)
from ..module import StateModule
from ..types import (
    JSONSerializableObject,
    ToolFunction,
)
from ..tracing._trace import trace_toolkit
from .._logging import logger


def _apply_middlewares(
    func: Callable[
        ...,
        Coroutine[Any, Any, AsyncGenerator[ToolResponse, None]],
    ],
) -> Callable[..., AsyncGenerator[ToolResponse, None]]:
    """Decorator that applies registered middlewares at runtime.

    This decorator reads the middleware list from the instance and constructs
    the middleware chain dynamically during each invocation.

    .. note:: Middlewares must be async generator functions that yield
     `ToolResponse` objects.
    """

    @wraps(func)
    async def wrapper(
        self: "Toolkit",
        tool_call: ToolUseBlock,
    ) -> AsyncGenerator[ToolResponse, None]:
        """Wrapper that applies middleware chain."""
        middlewares = getattr(self, "_middlewares", [])

        if not middlewares:
            # No middlewares, call the original function directly
            async for chunk in await func(self, tool_call):
                yield chunk
            return

        # Build the middleware chain from innermost to outermost
        async def base_handler(
            **kwargs: Any,
        ) -> AsyncGenerator[ToolResponse, None]:
            """Base handler that calls the original function."""
            return await func(self, **kwargs)

        # Wrap with each middleware in reverse order
        current_handler = base_handler
        for middleware in reversed(middlewares):

            def make_handler(mw: Callable, handler: Callable) -> Callable:
                """Create wrapped handler for middleware."""

                async def wrapped(
                    **kwargs: Any,
                ) -> AsyncGenerator[ToolResponse, None]:
                    """Handler that applies middleware."""
                    return mw(kwargs, handler)

                return wrapped

            current_handler = make_handler(middleware, current_handler)

        # Execute the middleware chain
        async for chunk in await current_handler(tool_call=tool_call):
            yield chunk

    return wrapper


class Toolkit(StateModule):
    """Toolkit is the core module to register, manage and delete tool
    functions, MCP clients, Agent skills in AgentScope.

    About tool functions:

    - Register and parse JSON schemas from their docstrings automatically.
    - Group-wise tools management, and agentic tools activation/deactivation.
    - Extend the tool function JSON schema dynamically with Pydantic BaseModel.
    - Tool function execution with unified streaming interface.

    About MCP clients:

    - Register tool functions from MCP clients directly.
    - Client-level tool functions removal.

    About Agent skills:

    - Register agent skills from the given directory.
    - Provide prompt for the registered skills to the agent.
    """

    _DEFAULT_AGENT_SKILL_INSTRUCTION = (
        "# Agent Skills\n"
        "The agent skills are a collection of folds of instructions, scripts, "
        "and resources that you can load dynamically to improve performance "
        "on specialized tasks. Each agent skill has a `SKILL.md` file in its "
        "folder that describes how to use the skill. If you want to use a "
        "skill, you MUST read its `SKILL.md` file carefully."
    )

    _DEFAULT_AGENT_SKILL_TEMPLATE = """## {name}
{description}
Check "{dir}/SKILL.md" for how to use this skill"""

    def __init__(
        self,
        agent_skill_instruction: str | None = None,
        agent_skill_template: str | None = None,
    ) -> None:
        """Initialize the toolkit.

        Args:
            agent_skill_instruction (`str | None`, optional):
                The instruction for agent skills in the system prompt. If not
                provided, a default instruction will be used.
            agent_skill_template (`str | None`, optional):
                The template to present one agent skill in the system prompt,
                which should contain `{name}`, `{description}`, and `{dir}`
                placeholders. If not provided, a default template will be used.
        """
        super().__init__()

        self.tools: dict[str, RegisteredToolFunction] = {}
        self.groups: dict[str, ToolGroup] = {}
        self.skills: dict[str, AgentSkill] = {}
        self._middlewares: list = []  # Store registered middlewares

        self._agent_skill_instruction = (
            agent_skill_instruction or self._DEFAULT_AGENT_SKILL_INSTRUCTION
        )
        self._agent_skill_template = (
            agent_skill_template or self._DEFAULT_AGENT_SKILL_TEMPLATE
        )

    def create_tool_group(
        self,
        group_name: str,
        description: str,
        active: bool = False,
        notes: str | None = None,
    ) -> None:
        """Create a tool group to organize tool functions

        Args:
            group_name (`str`):
                The name of the tool group.
            description (`str`):
                The description of the tool group.
            active (`bool`, defaults to `False`):
                If the group is active, meaning the tool functions in this
                group are included in the JSON schema.
            notes (`str | None`, optional):
                The notes used to remind the agent how to use the tool
                functions properly, which can be combined into the system
                prompt.
        """
        if group_name in self.groups or group_name == "basic":
            raise ValueError(
                f"Tool group '{group_name}' is already registered in the "
                "toolkit.",
            )

        self.groups[group_name] = ToolGroup(
            name=group_name,
            description=description,
            notes=notes,
            active=active,
        )

    def update_tool_groups(self, group_names: list[str], active: bool) -> None:
        """Update the activation status of the given tool groups.

        Args:
            group_names (`list[str]`):
                The list of tool group names to be updated.
            active (`bool`):
                If the tool groups should be activated or deactivated.
        """

        for group_name in group_names:
            if group_name == "basic":
                logger.warning(
                    "The 'basic' tool group is always active, skipping it.",
                )

            if group_name in self.groups:
                self.groups[group_name].active = active

    def remove_tool_groups(self, group_names: str | list[str]) -> None:
        """Remove tool functions from the toolkit by their group names.

        Args:
            group_names (`str | list[str]`):
                The group names to be removed from the toolkit.
        """
        if isinstance(group_names, str):
            group_names = [group_names]

        if not isinstance(group_names, list) or not all(
            isinstance(_, str) for _ in group_names
        ):
            raise TypeError(
                f"The group_names must be a list of strings, "
                f"but got {type(group_names)}.",
            )

        if "basic" in group_names:
            raise ValueError(
                "Cannot remove the default 'basic' tool group.",
            )

        for group_name in group_names:
            self.groups.pop(group_name, None)

        # Remove the tool functions in the given groups
        tool_names = deepcopy(list(self.tools.keys()))
        for tool_name in tool_names:
            if self.tools[tool_name].group in group_names:
                self.tools.pop(tool_name)

    # pylint: disable=too-many-branches, too-many-statements
    def register_tool_function(
        self,
        tool_func: ToolFunction,
        group_name: str | Literal["basic"] = "basic",
        preset_kwargs: dict[str, JSONSerializableObject] | None = None,
        func_name: str | None = None,
        func_description: str | None = None,
        json_schema: dict | None = None,
        include_long_description: bool = True,
        include_var_positional: bool = False,
        include_var_keyword: bool = False,
        postprocess_func: (
            Callable[
                [ToolUseBlock, ToolResponse],
                ToolResponse | None,
            ]
            | Callable[
                [ToolUseBlock, ToolResponse],
                Awaitable[ToolResponse | None],
            ]
        )
        | None = None,
        namesake_strategy: Literal[
            "override",
            "skip",
            "raise",
            "rename",
        ] = "raise",
    ) -> None:
        """Register a tool function to the toolkit.

        Args:
            tool_func (`ToolFunction`):
                The tool function, which can be async or sync, streaming or
                not-streaming, but the response must be a `ToolResponse`
                object.
            group_name (`str | Literal["basic"]`, defaults to `"basic"`):
                The belonging group of the tool function. Tools in "basic"
                group is always included in the JSON schema, while the others
                are only included when their group is active.
            preset_kwargs (`dict[str, JSONSerializableObject] | None`, \
            optional):
                Preset arguments by the user, which will not be included in
                the JSON schema, nor exposed to the agent.
            func_name (`str | None`, optional):
                The custom function name, which should be consistent with the
                name in function_description and json_schema (if provided).
                By default, the function name will be extracted from the
                function automatically.
            func_description (`str | None`, optional):
                The function description. If not provided, the description
                will be extracted from the docstring automatically.
            json_schema (`dict | None`, optional):
                Manually provided JSON schema for the tool function, which
                should be `{"type": "function", "function": {"name":
                "function_name": "xx", "description": "xx",
                "parameters": {...}}}`
            include_long_description (`bool`, defaults to `True`):
                When extracting function description from the docstring, if
                the long description will be included.
            include_var_positional (`bool`, defaults to `False`):
                Whether to include the variable positional arguments (`*args`)
                in the function schema.
            include_var_keyword (`bool`, defaults to `False`):
                Whether to include the variable keyword arguments (`**kwargs`)
                in the function schema.
            postprocess_func (`(Callable[[ToolUseBlock, ToolResponse], \
            ToolResponse | None] | Callable[[ToolUseBlock, ToolResponse], \
            Awaitable[ToolResponse | None]]) | None`, optional):
                A post-processing function that will be called after the tool
                function is executed, taking the tool call block and tool
                response as arguments. The function can be either sync or
                async. If it returns `None`, the tool result will be
                returned as is. If it returns a `ToolResponse`,
                the returned block will be used as the final tool result.
            namesake_strategy (`Literal['raise', 'override', 'skip', \
            'rename']`, defaults to `'raise'`):
                The strategy to handle the tool function name conflict:
                - 'raise': raise a ValueError (default behavior).
                - 'override': override the existing tool function with the new
                  one.
                - 'skip': skip the registration of the new tool function.
                - 'rename': rename the new tool function by appending a random
                  suffix to make it unique.
        """
        # Arguments checking
        if group_name not in self.groups and group_name != "basic":
            raise ValueError(
                f"Tool group '{group_name}' not found.",
            )

        # Check the manually provided JSON schema if provided
        if json_schema:
            assert (
                isinstance(json_schema, dict)
                and "type" in json_schema
                and json_schema["type"] == "function"
                and "function" in json_schema
                and isinstance(json_schema["function"], dict)
            ), "Invalid JSON schema for the tool function."

        # Handle MCP tool function and regular function respectively
        mcp_name = None
        if isinstance(tool_func, MCPToolFunction):
            input_func_name = tool_func.name
            original_func = tool_func.__call__
            json_schema = json_schema or tool_func.json_schema
            mcp_name = tool_func.mcp_name

        elif isinstance(tool_func, partial):
            # partial function
            kwargs = tool_func.keywords
            # Turn args into keyword arguments
            if tool_func.args:
                param_names = list(
                    inspect.signature(tool_func.func).parameters.keys(),
                )
                for i, arg in enumerate(tool_func.args):
                    if i < len(param_names):
                        kwargs[param_names[i]] = arg

            preset_kwargs = {
                **kwargs,
                **(preset_kwargs or {}),
            }

            input_func_name = tool_func.func.__name__
            original_func = tool_func.func
            json_schema = json_schema or _parse_tool_function(
                tool_func.func,
                include_long_description=include_long_description,
                include_var_positional=include_var_positional,
                include_var_keyword=include_var_keyword,
            )

        else:
            # normal function
            input_func_name = tool_func.__name__
            original_func = tool_func
            json_schema = json_schema or _parse_tool_function(
                tool_func,
                include_long_description=include_long_description,
                include_var_positional=include_var_positional,
                include_var_keyword=include_var_keyword,
            )

        # Record the original function name if the func_name is given
        original_name = input_func_name if func_name else None

        # Use the given function name if provided
        func_name = func_name or input_func_name

        # Always set the function name in json_schema
        json_schema["function"]["name"] = func_name

        # Override the description if provided
        if func_description:
            json_schema["function"]["description"] = func_description

        # Remove the preset kwargs from the JSON schema
        for arg_name in preset_kwargs or {}:
            if arg_name in json_schema["function"]["parameters"]["properties"]:
                json_schema["function"]["parameters"]["properties"].pop(
                    arg_name,
                )

        if "required" in json_schema["function"]["parameters"]:
            for arg_name in preset_kwargs or {}:
                if (
                    arg_name
                    in json_schema["function"]["parameters"]["required"]
                ):
                    json_schema["function"]["parameters"]["required"].remove(
                        arg_name,
                    )

            # Remove the required field if it is empty
            if len(json_schema["function"]["parameters"]["required"]) == 0:
                json_schema["function"]["parameters"].pop("required", None)

        func_obj = RegisteredToolFunction(
            name=func_name,
            group=group_name,
            source="function",
            original_func=original_func,
            json_schema=json_schema,
            preset_kwargs=preset_kwargs or {},
            original_name=original_name,
            extended_model=None,
            mcp_name=mcp_name,
            postprocess_func=postprocess_func,
        )

        if func_name in self.tools:
            if namesake_strategy == "raise":
                raise ValueError(
                    f"A function with name '{func_name}' is already "
                    f"registered in the toolkit.",
                )

            if namesake_strategy == "skip":
                logger.warning(
                    "A function with name '%s' is already "
                    "registered in the toolkit. Skipping registration.",
                    func_name,
                )

            elif namesake_strategy == "override":
                logger.warning(
                    "A function with name '%s' is already registered "
                    "in the toolkit. Overriding with the new function.",
                    func_name,
                )
                self.tools[func_name] = func_obj

            elif namesake_strategy == "rename":
                new_func_name = func_name
                for _ in range(100):
                    suffix = shortuuid.uuid()[:5]
                    new_func_name = f"{func_name}_{suffix}"
                    if new_func_name not in self.tools:
                        break

                # Raise error if failed to find a unique name
                if new_func_name in self.tools:
                    raise RuntimeError(
                        f"Failed to register tool function '{func_name}' with "
                        "a unique name after 100 attempts.",
                    )
                logger.warning(
                    "A function with name '%s' is already "
                    "registered in the toolkit. Renaming the new function to "
                    "'%s'.",
                    func_name,
                    new_func_name,
                )

                # Replace the function name with the new one
                func_obj.original_name = original_name or func_name
                func_obj.name = new_func_name
                func_obj.json_schema["function"]["name"] = new_func_name

                self.tools[new_func_name] = func_obj

            else:
                raise ValueError(
                    f"Invalid namesake_strategy: {namesake_strategy}. "
                    "Supported strategies are 'raise', 'override', 'skip', "
                    "and 'rename'.",
                )

        else:
            self.tools[func_name] = func_obj

    def remove_tool_function(
        self,
        tool_name: str,
        allow_not_exist: bool = True,
    ) -> None:
        """Remove tool function from the toolkit by its name.

        Args:
            tool_name (`str`):
                The name of the tool function to be removed.
            allow_not_exist (`bool`):
                Allow the tool function to not exist when removing.
        """

        if tool_name not in self.tools and not allow_not_exist:
            raise ValueError(
                f"Tool function '{tool_name}' does not exist in the "
                "toolkit.",
            )

        self.tools.pop(tool_name, None)

    def get_json_schemas(
        self,
    ) -> list[dict]:
        """Get the JSON schemas from the tool functions that belong to the
        active groups.

        .. note:: The preset keyword arguments is removed from the JSON
         schema, and the extended model is applied if it is set.

        Example:
            .. code-block:: JSON
                :caption: Example of tool function JSON schemas

                [
                    {
                        "type": "function",
                        "function": {
                            "name": "google_search",
                            "description": "Search on Google.",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "query": {
                                        "type": "string",
                                        "description": "The search query."
                                    }
                                },
                                "required": ["query"]
                            }
                        }
                    },
                    ...
                ]

        Returns:
            `list[dict]`:
                A list of function JSON schemas.
        """
        # If meta tool is set here, update its extended model here
        if "reset_equipped_tools" in self.tools:
            fields = {}
            for group_name, group in self.groups.items():
                if group_name == "basic":
                    continue
                fields[group_name] = (
                    bool,
                    Field(
                        default=False,
                        description=group.description,
                    ),
                )
            extended_model = create_model("_DynamicModel", **fields)
            self.set_extended_model(
                "reset_equipped_tools",
                extended_model,
            )

        return [
            tool.extended_json_schema
            for tool in self.tools.values()
            if tool.group == "basic" or self.groups[tool.group].active
        ]

    def set_extended_model(
        self,
        func_name: str,
        model: Type[BaseModel] | None,
    ) -> None:
        """Set the extended model for a tool function, so that the original
        JSON schema will be extended.

        Args:
            func_name (`str`):
                The name of the tool function.
            model (`Union[Type[BaseModel], None]`):
                The extended model to be set.
        """
        if model is not None and not issubclass(model, BaseModel):
            raise TypeError(
                "The extended model must be a child class of pydantic "
                f"BaseModel, but got {type(model)}.",
            )

        if func_name in self.tools:
            self.tools[func_name].extended_model = model

        else:
            raise ValueError(
                f"Tool function '{func_name}' not found in the toolkit.",
            )

    async def remove_mcp_clients(
        self,
        client_names: list[str],
    ) -> None:
        """Remove tool functions from the MCP clients by their names.

        Args:
            client_names (`list[str]`):
                The names of the MCP client, which used to initialize the
                client instance.
        """
        if isinstance(client_names, str):
            client_names = [client_names]

        if isinstance(client_names, list) and not all(
            isinstance(_, str) for _ in client_names
        ):
            raise TypeError(
                f"The client_names must be a list of strings, "
                f"but got {type(client_names)}.",
            )

        to_removed = []
        func_names = deepcopy(list(self.tools.keys()))
        for func_name in func_names:
            if self.tools[func_name].mcp_name in client_names:
                self.tools.pop(func_name)
                to_removed.append(func_name)

        logger.info(
            "Removed %d tool functions from %d MCP: %s",
            len(to_removed),
            len(client_names),
            ", ".join(to_removed),
        )

    @trace_toolkit
    @_apply_middlewares
    async def call_tool_function(
        self,
        tool_call: ToolUseBlock,
    ) -> AsyncGenerator[ToolResponse, None]:
        """Execute the tool function by the `ToolUseBlock` and return the
        tool response chunk in unified streaming mode, i.e. an async
        generator of `ToolResponse` objects.

        .. note:: The tool response chunk is **accumulated**.

        Args:
            tool_call (`ToolUseBlock`):
                A tool call block.

        Yields:
            `ToolResponse`:
                The tool response chunk, in accumulative manner.
        """

        # Check
        if tool_call["name"] not in self.tools:
            return _object_wrapper(
                ToolResponse(
                    content=[
                        TextBlock(
                            type="text",
                            text="FunctionNotFoundError: Cannot find the "
                            f"function named {tool_call['name']}",
                        ),
                    ],
                ),
                None,
            )

        # Obtain the tool function
        tool_func = self.tools[tool_call["name"]]

        # Check if the tool function is in an inactive group
        if (
            tool_func.group != "basic"
            and not self.groups[tool_func.group].active
        ):
            return _object_wrapper(
                ToolResponse(
                    content=[
                        TextBlock(
                            type="text",
                            text="FunctionInactiveError: The function "
                            f"'{tool_call['name']}' is in the inactive "
                            f"group '{tool_func.group}'. Activate the tool "
                            "group by calling 'reset_equipped_tools' "
                            "first to use this tool.",
                        ),
                    ],
                ),
                None,
            )

        # Prepare keyword arguments
        kwargs = {
            **tool_func.preset_kwargs,
            **(tool_call.get("input", {}) or {}),
        }

        # Prepare postprocess function
        if tool_func.postprocess_func:
            # Type: partial wraps the postprocess_func with tool_call bound,
            # reducing it from (ToolUseBlock, ToolResponse) to (ToolResponse)
            partial_postprocess_func: (
                Callable[[ToolResponse], ToolResponse | None]
                | Callable[[ToolResponse], Awaitable[ToolResponse | None]]
            ) | None = partial(
                tool_func.postprocess_func,
                tool_call,
            )
        else:
            partial_postprocess_func = None

        # Async function
        try:
            if inspect.iscoroutinefunction(tool_func.original_func):
                try:
                    res = await tool_func.original_func(**kwargs)
                except asyncio.CancelledError:
                    res = ToolResponse(
                        content=[
                            TextBlock(
                                type="text",
                                text="<system-info>"
                                "The tool call has been interrupted "
                                "by the user."
                                "</system-info>",
                            ),
                        ],
                        stream=True,
                        is_last=True,
                        is_interrupted=True,
                    )

            else:
                # When `tool_func.original_func` is Async generator function or
                # Sync function
                res = tool_func.original_func(**kwargs)

        except mcp.shared.exceptions.McpError as e:
            res = ToolResponse(
                content=[
                    TextBlock(
                        type="text",
                        text=f"Error occurred when calling MCP tool: {e}",
                    ),
                ],
            )

        except Exception as e:
            res = ToolResponse(
                content=[
                    TextBlock(
                        type="text",
                        text=f"Error: {e}",
                    ),
                ],
            )

        # Handle different return type

        # If return an async generator
        if isinstance(res, AsyncGenerator):
            return _async_generator_wrapper(res, partial_postprocess_func)

        # If return a sync generator
        if isinstance(res, Generator):
            return _sync_generator_wrapper(res, partial_postprocess_func)

        if isinstance(res, ToolResponse):
            return _object_wrapper(res, partial_postprocess_func)

        raise TypeError(
            "The tool function must return a ToolResponse object, or an "
            "AsyncGenerator/Generator of ToolResponse objects, "
            f"but got {type(res)}.",
        )

    async def register_mcp_client(
        self,
        mcp_client: MCPClientBase,
        group_name: str = "basic",
        enable_funcs: list[str] | None = None,
        disable_funcs: list[str] | None = None,
        preset_kwargs_mapping: dict[str, dict[str, Any]] | None = None,
        postprocess_func: (
            Callable[
                [ToolUseBlock, ToolResponse],
                ToolResponse | None,
            ]
            | Callable[
                [ToolUseBlock, ToolResponse],
                Awaitable[ToolResponse | None],
            ]
        )
        | None = None,
        namesake_strategy: Literal[
            "override",
            "skip",
            "raise",
            "rename",
        ] = "raise",
    ) -> None:
        """Register tool functions from an MCP client.

        Args:
            mcp_client (`MCPClientBase`):
                The MCP client instance to connect to the MCP server.
            group_name (`str`, defaults to `"basic"`):
                The group name that the tool functions will be added to.
            enable_funcs (`list[str] | None`, optional):
                The functions to be added into the toolkit. If `None`, all
                tool functions within the MCP servers will be added.
            disable_funcs (`list[str] | None`, optional):
                The functions that will be filtered out. If `None`, no
                tool functions will be filtered out.
            preset_kwargs_mapping: (`Optional[dict[str, dict[str, Any]]]`, \
            defaults to `None`):
                The preset keyword arguments mapping, whose keys are the tool
                function names and values are the preset keyword arguments.
            postprocess_func (`(Callable[[ToolUseBlock, ToolResponse], \
            ToolResponse | None] | Callable[[ToolUseBlock, ToolResponse], \
            Awaitable[ToolResponse | None]]) | None`, optional):
                A post-processing function that will be called after the tool
                function is executed, taking the tool call block and tool
                response as arguments. The function can be either sync or
                async. If it returns `None`, the tool result will be
                returned as is. If it returns a `ToolResponse`,
                the returned block will be used as the final tool result.
            namesake_strategy (`Literal['raise', 'override', 'skip', \
            'rename']`, defaults to `'raise'`):
                The strategy to handle the tool function name conflict:
                - 'raise': raise a ValueError (default behavior).
                - 'override': override the existing tool function with the new
                  one.
                - 'skip': skip the registration of the new tool function.
                - 'rename': rename the new tool function by appending a random
                  suffix to make it unique.
        """
        if (
            isinstance(mcp_client, StatefulClientBase)
            and not mcp_client.is_connected
        ):
            raise RuntimeError(
                "The MCP client is not connected to the server. Use the "
                "`connect()` method first.",
            )

        # Check arguments for enable_funcs and disabled_funcs
        if enable_funcs is not None and disable_funcs is not None:
            assert isinstance(enable_funcs, list) and all(
                isinstance(_, str) for _ in enable_funcs
            ), (
                "Enable functions should be a list of strings, but got "
                f"{enable_funcs}."
            )

            assert isinstance(disable_funcs, list) and all(
                isinstance(_, str) for _ in disable_funcs
            ), (
                "Disable functions should be a list of strings, but got "
                f"{disable_funcs}."
            )
            intersection = set(enable_funcs).intersection(
                set(disable_funcs),
            )
            assert len(intersection) == 0, (
                f"The functions in enable_funcs and disable_funcs "
                f"should not overlap, but got {intersection}."
            )

        if not (
            preset_kwargs_mapping is None
            or isinstance(preset_kwargs_mapping, dict)
        ):
            raise TypeError(
                f"The preset_kwargs_mapping must be a dictionary or None, "
                f"but got {type(preset_kwargs_mapping)}.",
            )

        tool_names = []
        for mcp_tool in await mcp_client.list_tools():
            # Skip the functions that are not in the enable_funcs if
            # enable_funcs is not None
            if enable_funcs is not None and mcp_tool.name not in enable_funcs:
                continue

            # Skip the disabled functions
            if disable_funcs is not None and mcp_tool.name in disable_funcs:
                continue

            tool_names.append(mcp_tool.name)

            # Obtain callable function object
            func_obj = await mcp_client.get_callable_function(
                func_name=mcp_tool.name,
                wrap_tool_result=True,
            )

            # Prepare preset kwargs
            preset_kwargs = None
            if preset_kwargs_mapping is not None:
                preset_kwargs = preset_kwargs_mapping.get(mcp_tool.name, {})

            self.register_tool_function(
                tool_func=func_obj,
                group_name=group_name,
                preset_kwargs=preset_kwargs,
                postprocess_func=postprocess_func,
                namesake_strategy=namesake_strategy,
            )

        logger.info(
            "Registered %d tool functions from MCP: %s.",
            len(tool_names),
            ", ".join(tool_names),
        )

    def state_dict(self) -> dict[str, Any]:
        """Get the state dictionary of the toolkit.

        Returns:
            `dict[str, Any]`:
                A dictionary containing the active tool group names.
        """
        return {
            "active_groups": [
                name for name, group in self.groups.items() if group.active
            ],
        }

    def load_state_dict(
        self,
        state_dict: dict[str, Any],
        strict: bool = True,
    ) -> None:
        """Load the state dictionary into the toolkit.

        Args:
            state_dict (`dict`):
                The state dictionary to load, which should have "active_groups"
                key and its value must be a list of group names.
            strict (`bool`, defaults to `True`):
                If `True`, raises an error if any key in the module is not
                found in the state_dict. If `False`, skips missing keys.
        """
        if (
            not isinstance(state_dict, dict)
            or "active_groups" not in state_dict
            or not isinstance(state_dict["active_groups"], list)
        ):
            raise ValueError(
                "The state_dict for toolkit must be a dictionary with "
                "active_groups key and its value must be a list, "
                f"but got {type(state_dict)}.",
            )

        if strict and list(state_dict.keys()) != ["active_groups"]:
            raise ValueError(
                "Get additional keys in the state_dict: "
                f'{list(state_dict.keys())}, but only "active_groups" '
                "is expected.",
            )

        for group_name, group in self.groups.items():
            if group_name in state_dict["active_groups"]:
                group.active = True
            else:
                group.active = False

    def get_activated_notes(self) -> str:
        """Get the notes from the active tool groups, which can be used to
        construct the system prompt for the agent.

        Returns:
            `str`:
                The combined notes from the active tool groups.
        """
        collected_notes = []
        for group_name, group in self.groups.items():
            if group.active and group.notes:
                collected_notes.append(
                    "\n".join(
                        [f"## About Tool Group '{group_name}'", group.notes],
                    ),
                )
        return "\n".join(collected_notes)

    def reset_equipped_tools(self, **kwargs: Any) -> ToolResponse:
        """This function allows you to activate or deactivate tool groups
        dynamically based on your current task requirements.
        **Important: Each call sets the absolute final state of ALL tool
        groups, not incremental changes**. Any group not explicitly set to True
        will be deactivated, regardless of its previous state.

        **Best practice**: Actively manage your tool groups——activate only
        what you need for the current task, and promptly deactivate groups as
        soon as they are no longer needed to conserve context space.

        The function will return the usage instructions for the activated tool
        groups, which you **MUST pay attention to and follow**. You can also
        reuse this function to check the notes of the tool groups."""

        # Deactivate all tool groups first
        self.update_tool_groups(list(self.groups.keys()), active=False)

        to_activate = []
        for key, value in kwargs.items():
            if not isinstance(value, bool):
                return ToolResponse(
                    content=[
                        TextBlock(
                            type="text",
                            text=f"Invalid arguments: the argument {key} "
                            f"should be a bool value, but got {type(value)}.",
                        ),
                    ],
                )

            if value:
                to_activate.append(key)

        self.update_tool_groups(to_activate, active=True)

        notes = self.get_activated_notes()

        text_response = ""
        if to_activate:
            text_response += (
                "Now tool groups "
                + ", ".join([f"'{_}'" for _ in to_activate])
                + " are activated."
            )

        if notes:
            text_response += (
                f" You MUST follow these notes to use these tools:\n"
                f"<notes>{notes}</notes>"
            )

        if not text_response:
            text_response = "All tool groups are now deactivated currently."

        return ToolResponse(
            content=[
                TextBlock(
                    type="text",
                    text=text_response,
                ),
            ],
        )

    def clear(self) -> None:
        """Clear the toolkit, removing all tool functions and groups."""
        self.tools.clear()
        self.groups.clear()

    def _validate_tool_function(self, func_name: str) -> None:
        """Check if the tool function already registered in the toolkit. If
        so, raise a ValueError."""
        if func_name in self.tools:
            raise ValueError(
                f"A function with name '{func_name}' is already registered "
                "in the toolkit.",
            )

    def register_agent_skill(
        self,
        skill_dir: str,
    ) -> None:
        """Register agent skills from a given directory. This function will
        scan the directory, read metadata from the SKILL.md file, and add
        it to the skill related prompt. Developers can obtain the
        skills-related prompt by calling `toolkit.get_agent_skill_prompt()`.

        .. note:: This directory
         - Must include a SKILL.md file at the top level
         - The SKILL.md must have a YAML Front Matter including `name` and
            `description` fields
         - All files must specify a common root directory in their paths

        Args:
            skill_dir (`str`):
                The path to the skill directory.
        """
        import frontmatter

        # Check the skill directory
        if not os.path.isdir(skill_dir):
            raise ValueError(
                f"The skill directory '{skill_dir}' does not exist or is "
                "not a directory.",
            )

        # Check SKILL.md file
        path_skill_md = os.path.join(skill_dir, "SKILL.md")
        if not os.path.isfile(path_skill_md):
            raise ValueError(
                f"The skill directory '{skill_dir}' must include a "
                "SKILL.md file at the top level.",
            )

        # Check YAML Front Matter
        with open(path_skill_md, "r", encoding="utf-8") as f:
            post = frontmatter.load(f)

        name = post.get("name", None)
        description = post.get("description", None)

        if not name or not description:
            raise ValueError(
                f"The SKILL.md file in '{skill_dir}' must have a YAML Front "
                "Matter including `name` and `description` fields.",
            )

        name, description = str(name), str(description)
        if name in self.skills:
            raise ValueError(
                f"An agent skill with name '{name}' is already registered "
                "in the toolkit.",
            )

        self.skills[name] = AgentSkill(
            name=name,
            description=description,
            dir=skill_dir,
        )

        logger.info(
            "Registered agent skill '%s' from directory '%s'.",
            name,
            skill_dir,
        )

    def remove_agent_skill(self, name: str) -> None:
        """Remove an agent skill by its name.

        Args:
            name (`str`):
                The name of the agent skill to be removed.
        """
        if name in self.skills:
            self.skills.pop(name)
        else:
            logger.warning(
                "Agent skill '%s' not found in the toolkit, skipping removal.",
                name,
            )

    def get_agent_skill_prompt(self) -> str | None:
        """Get the prompt for all registered agent skills, which can be
        attached to the system prompt for the agent.

        The prompt is consisted of an overall instruction and the detailed
        descriptions of each skill, including its name, description, and
        directory.

        .. note:: If no skill is registered, None will be returned.

        Returns:
            `str | None`:
                The combined prompt for all registered agent skills, or None
                if no skill is registered.
        """
        if len(self.skills) == 0:
            return None

        skill_descriptions = [
            self._agent_skill_instruction,
        ] + [
            self._agent_skill_template.format(
                name=_["name"],
                description=_["description"],
                dir=_["dir"],
            )
            for _ in self.skills.values()
        ]
        return "\n".join(skill_descriptions)

    def register_middleware(
        self,
        middleware: Callable[
            ...,
            Coroutine[Any, Any, AsyncGenerator[ToolResponse, None]]
            | AsyncGenerator[ToolResponse, None],
        ],
    ) -> None:
        """Register an onion-style middleware for the `call_tool_function`,
        which will wrap around the `call_tool_function` method, allowing
        pre-processing, post-processing, or even skipping the execution of
        the tool function.

        The middleware follows an onion model, where each registered
        middleware wraps around the previous one, forming layers. The
        middleware can:

        - Perform pre-processing before calling the tool function
        - Intercept and modify each ToolResponse chunk
        - Perform post-processing after the tool function completes
        - Skip the tool function execution entirely

        The middleware function should accept a ``kwargs`` dict as the first
        parameter and ``next_handler`` as the second parameter. The ``kwargs``
        dict currently contains:

        - ``tool_call`` (`ToolUseBlock`): The tool call request

        When calling ``next_handler``, pass ``**kwargs`` to unpack the dict.

        Example:
            .. code-block:: python

                # Simple direct consumption style (recommended)
                async def my_middleware(
                    kwargs: dict,
                    next_handler: Callable,
                ) -> AsyncGenerator[ToolResponse, None]:
                    # Access the tool call
                    tool_call = kwargs["tool_call"]

                    # Pre-processing
                    print(f"Calling tool: {tool_call['name']}")

                    # Call next handler with **kwargs
                    async for response in await next_handler(**kwargs):
                        # Intercept and modify response if needed
                        yield response

                    # Post-processing after tool completes
                    print(f"Tool {tool_call['name']} completed")

                toolkit.register_middleware(my_middleware)

            .. code-block:: python

                # Alternative: Skip execution based on conditions
                async def my_middleware(
                    kwargs: dict,
                    next_handler: Callable,
                ) -> AsyncGenerator[ToolResponse, None]:
                    tool_call = kwargs["tool_call"]

                    # Pre-processing
                    if not is_authorized(tool_call):
                        # Skip execution and return error directly
                        yield ToolResponse(
                            content=[
                                TextBlock(
                                    type="text",
                                    text="Unauthorized",
                                ),
                            ],
                        )
                        return

                    # Call next handler with **kwargs
                    async for response in await next_handler(**kwargs):
                        yield response

                toolkit.register_middleware(my_middleware)

        Args:
            middleware (`Callable[..., Coroutine[Any, Any, \
AsyncGenerator[ToolResponse, None]] | AsyncGenerator[ToolResponse, None]]`):
                The middleware function that accepts ``kwargs`` (dict) and
                ``next_handler`` (Callable), and returns a coroutine that
                yields AsyncGenerator of ToolResponse objects. The ``kwargs``
                dict currently includes ``tool_call`` (ToolUseBlock), and may
                include additional context in future versions.

        .. note:: The middleware chain is applied inside the
        `call_tool_function` via the `@apply_middlewares` decorator. This
        ensures that the `@trace_toolkit` decorator remains at the outermost
        layer for complete observability.
        """
        # Simply append the middleware to the list
        # The @apply_middlewares decorator will handle the execution
        self._middlewares.append(middleware)
