"""集成测试 - ReActEngine 并行工具执行

测试 parallel_tools 配置下的并行/串行/混合执行模式。
仅 mock LLMGateway（外部 API），使用真实 ReActEngine 实例。
"""

from __future__ import annotations

import asyncio
import time
from unittest.mock import AsyncMock, MagicMock

import pytest

from agentkit.core.react import ReActEngine, ReActResult, ReActEvent
from agentkit.llm.gateway import LLMGateway
from agentkit.llm.protocol import LLMResponse, TokenUsage, ToolCall
from agentkit.tools.base import Tool


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


class SlowTool(Tool):
    """带延迟的 Fake Tool，用于验证并行执行"""

    def __init__(
        self,
        name: str = "slow_tool",
        description: str = "A slow tool for testing parallel execution",
        delay: float = 0.1,
        result: dict | None = None,
    ):
        super().__init__(name=name, description=description)
        self._delay = delay
        self._result = result or {"status": "ok"}
        self.call_count = 0
        self.call_times: list[float] = []

    async def execute(self, **kwargs) -> dict:
        self.call_count += 1
        self.call_times.append(time.monotonic())
        await asyncio.sleep(self._delay)
        return self._result


def make_response(
    content: str = "",
    tool_calls: list[ToolCall] | None = None,
    prompt_tokens: int = 10,
    completion_tokens: int = 20,
) -> LLMResponse:
    return LLMResponse(
        content=content,
        model="test-model",
        usage=TokenUsage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
        ),
        tool_calls=tool_calls or [],
    )


def make_mock_gateway(responses: list[LLMResponse]) -> MagicMock:
    gateway = MagicMock(spec=LLMGateway)
    gateway.chat = AsyncMock(side_effect=responses)
    return gateway


# ---------------------------------------------------------------------------
# Test 1: parallel_tools=True — 所有工具并行执行
# ---------------------------------------------------------------------------


class TestParallelToolsTrue:
    """parallel_tools=True 时所有工具并行执行"""

    @pytest.mark.asyncio
    async def test_all_tools_execute_in_parallel(self):
        """并行执行应比串行更快"""
        tool_a = SlowTool(name="tool_a", delay=0.15)
        tool_b = SlowTool(name="tool_b", delay=0.15)

        # LLM 返回 2 个 tool_calls
        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1}),
                ToolCall(id="tc_2", name="tool_b", arguments={"y": 2}),
            ],
        )
        final_response = make_response(content="Done")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools=True)

        start = time.monotonic()
        result = await engine.execute(
            messages=[{"role": "user", "content": "Run both tools"}],
            tools=[tool_a, tool_b],
        )
        elapsed = time.monotonic() - start

        assert isinstance(result, ReActResult)
        assert tool_a.call_count == 1
        assert tool_b.call_count == 1
        # 并行执行：总时间应接近单个工具延迟，而非两倍
        assert elapsed < 0.35, f"Parallel execution too slow: {elapsed:.2f}s"

    @pytest.mark.asyncio
    async def test_parallel_result_in_trajectory(self):
        """并行执行的工具结果正确记录在 trajectory 中"""
        tool_a = SlowTool(name="tool_a", delay=0.05, result={"value": "A"})
        tool_b = SlowTool(name="tool_b", delay=0.05, result={"value": "B"})

        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1}),
                ToolCall(id="tc_2", name="tool_b", arguments={"y": 2}),
            ],
        )
        final_response = make_response(content="Combined result")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools=True)

        result = await engine.execute(
            messages=[{"role": "user", "content": "Run both"}],
            tools=[tool_a, tool_b],
        )

        tool_steps = [s for s in result.trajectory if s.action == "tool_call"]
        assert len(tool_steps) == 2
        names = {s.tool_name for s in tool_steps}
        assert names == {"tool_a", "tool_b"}


# ---------------------------------------------------------------------------
# Test 2: parallel_tools="auto" + LLM 标记 _parallelizable=true → 并行
# ---------------------------------------------------------------------------


class TestParallelToolsAuto:
    """parallel_tools="auto" 时根据 _parallelizable 标记决定并行"""

    @pytest.mark.asyncio
    async def test_auto_parallelizable_tools_execute_in_parallel(self):
        """auto 模式下 _parallelizable=true 的工具并行执行"""
        tool_a = SlowTool(name="tool_a", delay=0.15)
        tool_b = SlowTool(name="tool_b", delay=0.15)

        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1, "_parallelizable": True}),
                ToolCall(id="tc_2", name="tool_b", arguments={"y": 2, "_parallelizable": True}),
            ],
        )
        final_response = make_response(content="Auto parallel done")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools="auto")

        start = time.monotonic()
        result = await engine.execute(
            messages=[{"role": "user", "content": "Run in parallel"}],
            tools=[tool_a, tool_b],
        )
        elapsed = time.monotonic() - start

        assert isinstance(result, ReActResult)
        assert tool_a.call_count == 1
        assert tool_b.call_count == 1
        # 并行执行应比串行快
        assert elapsed < 0.35, f"Auto parallel execution too slow: {elapsed:.2f}s"

    @pytest.mark.asyncio
    async def test_auto_mixed_parallelizable_serial_then_parallel(self):
        """auto 模式下混合标记：先串行执行非并行工具，再并行执行并行工具"""
        tool_serial = SlowTool(name="tool_serial", delay=0.1)
        tool_para_a = SlowTool(name="tool_para_a", delay=0.1)
        tool_para_b = SlowTool(name="tool_para_b", delay=0.1)

        tool_call_response = make_response(
            content="",
            tool_calls=[
                # 串行工具
                ToolCall(id="tc_1", name="tool_serial", arguments={"x": 1}),
                # 并行工具
                ToolCall(id="tc_2", name="tool_para_a", arguments={"y": 2, "_parallelizable": True}),
                ToolCall(id="tc_3", name="tool_para_b", arguments={"z": 3, "_parallelizable": True}),
            ],
        )
        final_response = make_response(content="Mixed result")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools="auto")

        result = await engine.execute(
            messages=[{"role": "user", "content": "Mixed execution"}],
            tools=[tool_serial, tool_para_a, tool_para_b],
        )

        assert isinstance(result, ReActResult)
        assert tool_serial.call_count == 1
        assert tool_para_a.call_count == 1
        assert tool_para_b.call_count == 1

        # 验证串行工具先执行，并行工具后执行
        # 串行工具的调用时间应早于并行工具
        assert tool_serial.call_times[0] < tool_para_a.call_times[0]

        # 并行工具的调用时间应接近（几乎同时开始）
        para_diff = abs(tool_para_a.call_times[0] - tool_para_b.call_times[0])
        assert para_diff < 0.05, f"Parallel tools didn't start together: diff={para_diff:.3f}s"

    @pytest.mark.asyncio
    async def test_auto_no_parallelizable_tools_execute_serially(self):
        """auto 模式下没有 _parallelizable 标记时串行执行"""
        tool_a = SlowTool(name="tool_a", delay=0.1)
        tool_b = SlowTool(name="tool_b", delay=0.1)

        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1}),
                ToolCall(id="tc_2", name="tool_b", arguments={"y": 2}),
            ],
        )
        final_response = make_response(content="Serial done")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools="auto")

        start = time.monotonic()
        result = await engine.execute(
            messages=[{"role": "user", "content": "Run serially"}],
            tools=[tool_a, tool_b],
        )
        elapsed = time.monotonic() - start

        assert isinstance(result, ReActResult)
        assert tool_a.call_count == 1
        assert tool_b.call_count == 1
        # 串行执行：总时间应接近两倍延迟
        assert elapsed >= 0.15, f"Serial execution unexpectedly fast: {elapsed:.2f}s"


# ---------------------------------------------------------------------------
# Test 3: parallel_tools=False — 全部串行（默认行为）
# ---------------------------------------------------------------------------


class TestParallelToolsFalse:
    """parallel_tools=False 时所有工具串行执行"""

    @pytest.mark.asyncio
    async def test_all_tools_execute_serially(self):
        """默认串行执行"""
        tool_a = SlowTool(name="tool_a", delay=0.1)
        tool_b = SlowTool(name="tool_b", delay=0.1)

        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1}),
                ToolCall(id="tc_2", name="tool_b", arguments={"y": 2}),
            ],
        )
        final_response = make_response(content="Serial result")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools=False)

        start = time.monotonic()
        result = await engine.execute(
            messages=[{"role": "user", "content": "Run serially"}],
            tools=[tool_a, tool_b],
        )
        elapsed = time.monotonic() - start

        assert isinstance(result, ReActResult)
        assert tool_a.call_count == 1
        assert tool_b.call_count == 1
        # 串行执行：总时间应接近两倍延迟
        assert elapsed >= 0.15, f"Serial execution unexpectedly fast: {elapsed:.2f}s"

    @pytest.mark.asyncio
    async def test_default_is_serial(self):
        """默认 parallel_tools=False"""
        tool_a = SlowTool(name="tool_a", delay=0.1)
        tool_b = SlowTool(name="tool_b", delay=0.1)

        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1}),
                ToolCall(id="tc_2", name="tool_b", arguments={"y": 2}),
            ],
        )
        final_response = make_response(content="Default serial")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway)  # default parallel_tools=False

        start = time.monotonic()
        result = await engine.execute(
            messages=[{"role": "user", "content": "Default mode"}],
            tools=[tool_a, tool_b],
        )
        elapsed = time.monotonic() - start

        assert isinstance(result, ReActResult)
        assert elapsed >= 0.15, f"Default mode should be serial, but was fast: {elapsed:.2f}s"

    @pytest.mark.asyncio
    async def test_single_tool_call_is_serial(self):
        """单个工具调用无论 parallel_tools 设置如何都是串行"""
        tool_a = SlowTool(name="tool_a", delay=0.05)

        tool_call_response = make_response(
            content="",
            tool_calls=[
                ToolCall(id="tc_1", name="tool_a", arguments={"x": 1}),
            ],
        )
        final_response = make_response(content="Single tool done")

        gateway = make_mock_gateway([tool_call_response, final_response])
        engine = ReActEngine(llm_gateway=gateway, parallel_tools=True)

        result = await engine.execute(
            messages=[{"role": "user", "content": "Single tool"}],
            tools=[tool_a],
        )

        assert isinstance(result, ReActResult)
        assert tool_a.call_count == 1
