# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

from ag2 import Agent, Context, tool
from ag2.events import ToolCallEvent
from ag2.exceptions import ToolConflictError
from ag2.middleware import ToolExecution, ToolResultType
from ag2.testing import TestConfig
from ag2.tools import Toolkit


@pytest.mark.asyncio
async def test_toolkit_schemas(async_mock: AsyncMock) -> None:
    @tool
    def add(a: int, b: int) -> int:
        """Add two numbers."""
        return a + b

    @tool
    def multiply(a: int, b: int) -> int:
        """Multiply two numbers."""
        return a * b

    toolkit = Toolkit(add, multiply)
    schemas = list(await toolkit.schemas(Context(async_mock)))

    assert len(schemas) == 2
    assert schemas[0].function.name == "add"
    assert schemas[1].function.name == "multiply"


@pytest.mark.asyncio()
async def test_toolkit_executes_tool(mock: MagicMock) -> None:
    @tool
    def add(a: int, b: int) -> int:
        """Add two numbers."""
        mock(a=a, b=b)
        return a + b

    toolkit = Toolkit(add)

    config = TestConfig(
        ToolCallEvent(name="add", arguments=json.dumps({"a": 2, "b": 3})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    result = await agent.ask("Hi!")

    mock.assert_called_once_with(a=2, b=3)
    assert result.body == "done"


@pytest.mark.asyncio()
async def test_toolkit_multiple_tools(mock: MagicMock) -> None:
    @tool
    def add(a: int, b: int) -> int:
        """Add two numbers."""
        mock.add(a=a, b=b)
        return a + b

    @tool
    def multiply(a: int, b: int) -> int:
        """Multiply two numbers."""
        mock.multiply(a=a, b=b)
        return a * b

    toolkit = Toolkit(add, multiply)

    config = TestConfig(
        ToolCallEvent(name="multiply", arguments=json.dumps({"a": 4, "b": 5})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    mock.add.assert_not_called()
    mock.multiply.assert_called_once_with(a=4, b=5)


@pytest.mark.asyncio()
async def test_toolkit_mixed_with_standalone_tool(mock: MagicMock) -> None:
    @tool
    def bundled(a: str) -> str:
        """Bundled tool."""
        mock.bundled(a)
        return a

    @tool
    def standalone(b: str) -> str:
        """Standalone tool."""
        mock.standalone(b)
        return b

    toolkit = Toolkit(bundled)

    config = TestConfig(
        ToolCallEvent(name="standalone", arguments=json.dumps({"b": "hello"})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit, standalone])
    await agent.ask("Hi!")

    mock.bundled.assert_not_called()
    mock.standalone.assert_called_once_with("hello")


@pytest.mark.asyncio()
async def test_toolkit_with_context(mock: MagicMock) -> None:
    from ag2 import Context

    @tool
    def greet(name: str, ctx: Context) -> str:
        """Greet someone."""
        mock(ctx.dependencies["lang"])
        return f"hello {name}"

    toolkit = Toolkit(greet)

    config = TestConfig(
        ToolCallEvent(name="greet", arguments=json.dumps({"name": "world"})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit], dependencies={"lang": "en"})
    await agent.ask("Hi!")

    mock.assert_called_once_with("en")


@pytest.mark.asyncio()
async def test_toolkit_with_plain_functions(mock: MagicMock) -> None:
    def add(a: int, b: int) -> int:
        """Add two numbers."""
        mock(a=a, b=b)
        return a + b

    toolkit = Toolkit(add)

    config = TestConfig(
        ToolCallEvent(name="add", arguments=json.dumps({"a": 1, "b": 2})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    mock.assert_called_once_with(a=1, b=2)


@pytest.mark.asyncio()
async def test_toolkit_mixed_functions_and_tools(mock: MagicMock) -> None:
    @tool
    def decorated(a: str) -> str:
        """Decorated tool."""
        mock.decorated(a)
        return a

    def plain(b: str) -> str:
        """Plain function."""
        mock.plain(b)
        return b

    toolkit = Toolkit(decorated, plain)

    config = TestConfig(
        ToolCallEvent(name="plain", arguments=json.dumps({"b": "hi"})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    mock.decorated.assert_not_called()
    mock.plain.assert_called_once_with("hi")


@pytest.mark.asyncio()
async def test_toolkit_tool_decorator(mock: MagicMock) -> None:
    toolkit = Toolkit()

    @toolkit.tool
    def greet(name: str) -> str:
        """Greet someone."""
        mock(name)
        return f"hello {name}"

    config = TestConfig(
        ToolCallEvent(name="greet", arguments=json.dumps({"name": "world"})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    mock.assert_called_once_with("world")


@pytest.mark.asyncio()
async def test_toolkit_tool_decorator_with_options(mock: MagicMock) -> None:
    toolkit = Toolkit()

    @toolkit.tool(name="say_hi", description="Custom greeting.")
    def greet(name: str) -> str:
        """Greet someone."""
        mock(name)
        return f"hello {name}"

    config = TestConfig(
        ToolCallEvent(name="say_hi", arguments=json.dumps({"name": "world"})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    mock.assert_called_once_with("world")


@pytest.mark.asyncio()
async def test_toolkit_empty() -> None:
    toolkit = Toolkit()

    config = TestConfig("done")
    agent = Agent("", config=config, tools=[toolkit])
    result = await agent.ask("Hi!")

    assert result.body == "done"


@pytest.mark.asyncio()
async def test_toolkit_middleware_applied_to_all_tools(mock: MagicMock) -> None:
    """Toolkit middleware wraps every tool in the set."""

    async def logging_middleware(
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        mock.before(event.name)
        result = await call_next(event, context)
        mock.after(event.name)
        return result

    @tool
    def add(a: int, b: int) -> int:
        """Add two numbers."""
        return a + b

    @tool
    def multiply(a: int, b: int) -> int:
        """Multiply two numbers."""
        return a * b

    toolkit = Toolkit(add, multiply, middleware=[logging_middleware])

    config = TestConfig(
        ToolCallEvent(name="add", arguments=json.dumps({"a": 1, "b": 2})),
        ToolCallEvent(name="multiply", arguments=json.dumps({"a": 3, "b": 4})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    assert mock.before.call_count == 2
    mock.before.assert_any_call("add")
    mock.before.assert_any_call("multiply")
    assert mock.after.call_count == 2


@pytest.mark.asyncio()
async def test_toolkit_middleware_wraps_subagent_tool(mock: MagicMock) -> None:
    async def logging_middleware(
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        mock(event.name)
        return await call_next(event, context)

    child = Agent("child", config=TestConfig("child answer"))
    toolkit = Toolkit(child.as_tool(description="Delegate work"), middleware=[logging_middleware])
    parent = Agent(
        "parent",
        config=TestConfig(
            ToolCallEvent(name="task_child", arguments=json.dumps({"objective": "work"})),
            "done",
        ),
        tools=[toolkit],
    )

    assert (await parent.ask("Hi!")).body == "done"
    mock.assert_called_once_with("task_child")


@pytest.mark.asyncio()
async def test_toolkit_middleware_applied_to_decorator_tools(mock: MagicMock) -> None:
    """Toolkit middleware also wraps tools added via the .tool() decorator."""

    async def logging_middleware(
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        mock.before(event.name)
        result = await call_next(event, context)
        mock.after(event.name)
        return result

    toolkit = Toolkit(middleware=[logging_middleware])

    @toolkit.tool
    def greet(name: str) -> str:
        """Greet someone."""
        return f"hello {name}"

    config = TestConfig(
        ToolCallEvent(name="greet", arguments=json.dumps({"name": "world"})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    mock.before.assert_called_once_with("greet")
    mock.after.assert_called_once_with("greet")


@pytest.mark.asyncio()
async def test_toolkit_middleware_ordering() -> None:
    """Per-tool middleware runs before toolkit middleware."""
    call_order: list[str] = []

    async def tool_mw(
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        call_order.append("tool_mw")
        return await call_next(event, context)

    async def toolkit_mw(
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        call_order.append("toolkit_mw")
        return await call_next(event, context)

    @tool(middleware=[tool_mw])
    def add(a: int, b: int) -> int:
        """Add two numbers."""
        call_order.append("tool")
        return a + b

    toolkit = Toolkit(add, middleware=[toolkit_mw])

    config = TestConfig(
        ToolCallEvent(name="add", arguments=json.dumps({"a": 1, "b": 2})),
        "done",
    )
    agent = Agent("", config=config, tools=[toolkit])
    await agent.ask("Hi!")

    assert call_order == ["toolkit_mw", "tool_mw", "tool"]


def test_tool_name_conflict() -> None:
    def add(a: int, b: int) -> int:
        pass

    with pytest.raises(ToolConflictError, match="add"):
        Toolkit(add, add)

    toolkit = Toolkit(add)
    with pytest.raises(ToolConflictError, match="add"):
        toolkit.tool(add)


class TestMerger:
    def test_merge_toolkits(self) -> None:
        left = Toolkit(first_tool)
        right = Toolkit(second_tool)
        members = (*left.tools, *right.tools)

        merged = left | right

        assert merged.tools == members

    def test_merge_toolkit_and_tool(self) -> None:
        def add1(a: int, b: int) -> int:
            pass

        def add2(a: int, b: int) -> int:
            pass

        toolkit = Toolkit(add1) | add2

        assert [t.name for t in toolkit.tools] == ["add1", "add2"]


async def first_tool(context: Context) -> str:
    context.dependencies["calls"].append("body:first")
    return "first"


async def second_tool(context: Context) -> str:
    context.dependencies["calls"].append("body:second")
    return "second"


async def third_tool(context: Context) -> str:
    context.dependencies["calls"].append("body:third")
    return "third"


async def left_middleware(call_next: ToolExecution, event: ToolCallEvent, context: Context) -> ToolResultType:
    context.dependencies["calls"].append(f"left:{event.name}")
    return await call_next(event, context)


async def right_middleware(call_next: ToolExecution, event: ToolCallEvent, context: Context) -> ToolResultType:
    context.dependencies["calls"].append(f"right:{event.name}")
    return await call_next(event, context)


async def tool_middleware(call_next: ToolExecution, event: ToolCallEvent, context: Context) -> ToolResultType:
    context.dependencies["calls"].append(f"tool:{event.name}")
    return await call_next(event, context)


async def old_tool_middleware(call_next: ToolExecution, event: ToolCallEvent, context: Context) -> ToolResultType:
    context.dependencies["calls"].append(f"old:{event.name}")
    return await call_next(event, context)


async def run_toolkit(toolkit: Toolkit, tool_name: str) -> list[str]:
    calls: list[str] = []
    agent = Agent(
        "test",
        tools=[toolkit],
        dependencies={"calls": calls},
        config=TestConfig(ToolCallEvent(name=tool_name), "done"),
    )
    await agent.ask("Run the tools")
    return calls


@pytest.mark.asyncio
class TestMergeMiddleware:
    async def test_repeated_merges_do_not_repeat_retained_middleware(self) -> None:
        left = Toolkit(first_tool, middleware=[left_middleware])
        right = Toolkit(second_tool, middleware=[right_middleware])

        merged = (left | right) | third_tool

        assert await run_toolkit(merged, "first_tool") == ["left:first_tool", "body:first"]

    async def test_incoming_tool_keeps_its_own_middleware(self) -> None:
        left = Toolkit(middleware=[left_middleware])
        right = Toolkit(tool(second_tool, middleware=[tool_middleware]), middleware=[right_middleware])

        merged = left | right

        assert await run_toolkit(merged, "second_tool") == [
            "right:second_tool",
            "tool:second_tool",
            "body:second",
        ]

    async def test_merge_leaves_source_toolkits_unchanged(self) -> None:
        left = Toolkit(first_tool, middleware=[left_middleware])
        right = Toolkit(second_tool, middleware=[right_middleware])

        merged = left | right
        merged.tool(third_tool)

        assert [t.name for t in left.tools] == ["first_tool"]
        assert [t.name for t in right.tools] == ["second_tool"]
        assert await run_toolkit(left, "first_tool") == ["left:first_tool", "body:first"]
        assert await run_toolkit(right, "second_tool") == ["right:second_tool", "body:second"]

    async def test_tool_added_after_merge_does_not_inherit_source_middleware(self) -> None:
        left = Toolkit(middleware=[left_middleware])
        merged = left | second_tool

        merged.tool()(third_tool)

        assert await run_toolkit(merged, "third_tool") == ["body:third"]

    async def test_tool_replaces_same_named_tool_and_its_hooks(self) -> None:
        old = tool(first_tool, name="shared", middleware=[old_tool_middleware])
        left = Toolkit(old)
        right = tool(second_tool, name="shared", middleware=[tool_middleware])

        merged = left | right

        assert await run_toolkit(merged, "shared") == ["tool:shared", "body:second"]