"""集成测试 - ReWOOEngine 可配置回退链

测试 fallback_strategies 参数控制下的自定义回退顺序。
仅 mock LLMGateway（外部 API），使用真实 ReWOOEngine 实例。
"""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

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


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


class FakeTool(Tool):
    """用于测试的 Fake Tool"""

    def __init__(
        self,
        name: str = "fake_tool",
        description: str = "A fake tool for testing",
        result: dict | None = None,
        should_fail: bool = False,
    ):
        super().__init__(name=name, description=description)
        self._result = result or {"status": "ok"}
        self._should_fail = should_fail
        self.call_count = 0

    async def execute(self, **kwargs) -> dict:
        self.call_count += 1
        if self._should_fail:
            raise RuntimeError(f"Tool '{self.name}' execution failed")
        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_plan_response(
    steps: list[dict],
    reasoning: str = "Plan reasoning",
) -> LLMResponse:
    plan_json = json.dumps({
        "reasoning": reasoning,
        "steps": steps,
    })
    return make_response(content=plan_json)


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


# ---------------------------------------------------------------------------
# Test 1: Custom fallback ["plan_exec", "react"] — planning 失败，plan_exec 失败，react 成功
# ---------------------------------------------------------------------------


class TestCustomFallbackPlanExecReact:
    """自定义回退链 ["plan_exec", "react"]"""

    @pytest.mark.asyncio
    async def test_plan_exec_fails_react_succeeds(self):
        """planning 失败 → plan_exec 失败 → react 成功"""
        tool = FakeTool(name="search", result={"results": ["found"]})

        # 1st: planning fails (invalid JSON)
        # 2nd: plan_exec planning also fails
        # 3rd: ReAct succeeds (tool call + final answer)
        invalid_plan = make_response(content="Cannot create a plan")
        invalid_plan_exec = make_response(content="Also cannot plan")
        react_tool_call = make_response(
            content="",
            tool_calls=[ToolCall(id="tc_1", name="search", arguments={"query": "test"})],
        )
        react_final = make_response(content="ReAct fallback answer")

        gateway = make_mock_gateway([
            invalid_plan,
            invalid_plan_exec,
            react_tool_call,
            react_final,
        ])
        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["plan_exec", "react"],
        )

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

        assert isinstance(result, ReActResult)
        assert result.output == "ReAct fallback answer"
        assert result.fallback_strategy == "react"
        assert tool.call_count == 1

    @pytest.mark.asyncio
    async def test_plan_exec_succeeds(self):
        """planning 失败 → plan_exec 成功"""
        tool = FakeTool(name="search", result={"results": ["found"]})

        # 1st: planning fails
        # 2nd: plan_exec planning succeeds
        # 3rd: synthesis
        invalid_plan = make_response(content="Cannot create a plan")
        plan_exec_response = make_plan_response([
            {"step_id": 1, "tool_name": "search", "arguments": {"query": "test"}, "reasoning": "Search"},
        ])
        synthesis = make_response(content="Plan-exec result")

        gateway = make_mock_gateway([invalid_plan, plan_exec_response, synthesis])
        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["plan_exec", "react"],
        )

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

        assert result.output == "Plan-exec result"
        assert result.fallback_strategy == "plan_exec"
        assert tool.call_count == 1


# ---------------------------------------------------------------------------
# Test 2: Custom fallback ["direct"] — planning 失败，direct 成功
# ---------------------------------------------------------------------------


class TestCustomFallbackDirect:
    """自定义回退链 ["direct"]"""

    @pytest.mark.asyncio
    async def test_direct_fallback_succeeds(self):
        """planning 失败 → direct LLM 调用成功"""
        invalid_plan = make_response(content="Cannot plan this task")
        direct_response = make_response(content="Direct LLM answer")

        gateway = make_mock_gateway([invalid_plan, direct_response])
        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["direct"],
        )

        result = await engine.execute(
            messages=[{"role": "user", "content": "Complex task"}],
        )

        assert isinstance(result, ReActResult)
        assert result.output == "Direct LLM answer"
        assert result.fallback_strategy == "direct"

    @pytest.mark.asyncio
    async def test_direct_fallback_with_system_prompt(self):
        """direct 回退时使用 system_prompt"""
        invalid_plan = make_response(content="Cannot plan")
        direct_response = make_response(content="Answer with context")

        gateway = make_mock_gateway([invalid_plan, direct_response])
        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["direct"],
        )

        result = await engine.execute(
            messages=[{"role": "user", "content": "Complex task"}],
            system_prompt="You are a helpful assistant",
        )

        assert result.output == "Answer with context"
        assert result.fallback_strategy == "direct"


# ---------------------------------------------------------------------------
# Test 3: Default fallback — 所有策略失败 → RuntimeError
# ---------------------------------------------------------------------------


class TestDefaultFallbackAllFail:
    """默认回退链所有策略都失败时抛出 RuntimeError"""

    @pytest.mark.asyncio
    async def test_all_fallback_strategies_exhausted(self):
        """所有回退策略耗尽 → RuntimeError"""
        # 所有 LLM 调用都返回无效 JSON
        # 默认 fallback: simplified_rewoo → react → direct
        # simplified_rewoo: 1 invalid plan
        # react: 1 response (but we make it fail by having the gateway raise)
        # direct: 1 response (but we make it fail too)

        # Use a gateway that raises on every call after the first
        call_count = 0

        class FailingGateway:
            def __init__(self):
                self.chat = AsyncMock(side_effect=self._fail_after_first)

            async def _fail_after_first(self, **kwargs):
                nonlocal call_count
                call_count += 1
                if call_count == 1:
                    # First call: planning fails (invalid JSON)
                    return make_response(content="Not a plan")
                elif call_count == 2:
                    # Second call: simplified_rewoo planning also fails
                    return make_response(content="Still not a plan")
                else:
                    # All subsequent calls fail
                    raise RuntimeError("LLM service unavailable")

        gateway = FailingGateway()
        engine = ReWOOEngine(llm_gateway=gateway)

        with pytest.raises(RuntimeError, match="All ReWOO fallback strategies exhausted"):
            await engine.execute(
                messages=[{"role": "user", "content": "Impossible task"}],
            )

    @pytest.mark.asyncio
    async def test_default_fallback_strategies_order(self):
        """验证默认回退策略顺序为 simplified_rewoo → react → direct"""
        engine = ReWOOEngine(
            llm_gateway=MagicMock(spec=LLMGateway),
        )
        assert engine._fallback_strategies == ["simplified_rewoo", "react", "direct"]


# ---------------------------------------------------------------------------
# Test 4: Stream mode with custom fallback
# ---------------------------------------------------------------------------


class TestStreamModeCustomFallback:
    """流式模式下的自定义回退链"""

    @pytest.mark.asyncio
    async def test_stream_direct_fallback(self):
        """流式模式下 planning 失败 → direct 回退"""
        invalid_plan = make_response(content="Cannot plan")
        direct_response = make_response(content="Stream direct answer")

        gateway = make_mock_gateway([invalid_plan, direct_response])
        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["direct"],
        )

        events = []
        async for event in engine.execute_stream(
            messages=[{"role": "user", "content": "Stream task"}],
        ):
            events.append(event)

        # 应该有 planning 事件和 final_answer 事件
        event_types = [e.event_type for e in events]
        assert "planning" in event_types
        assert "final_answer" in event_types

        # 最终答案来自 direct 回退
        final_events = [e for e in events if e.event_type == "final_answer"]
        assert len(final_events) >= 1
        assert final_events[0].data["output"] == "Stream direct answer"

    @pytest.mark.asyncio
    async def test_stream_react_fallback(self):
        """流式模式下 planning 失败 → react 回退"""
        tool = FakeTool(name="search", result={"results": ["found"]})

        invalid_plan = make_response(content="Cannot plan")

        # ReAct stream mode needs chat_stream as an async generator
        from agentkit.llm.protocol import StreamChunk

        async def mock_chat_stream(**kwargs):
            # First chunk: tool call
            yield StreamChunk(
                content="",
                model="test-model",
                tool_calls=[ToolCall(id="tc_1", name="search", arguments={"query": "test"})],
                usage=TokenUsage(prompt_tokens=10, completion_tokens=10),
            )
            # Second chunk: final answer
            yield StreamChunk(
                content="Stream react answer",
                model="test-model",
                tool_calls=[],
                usage=TokenUsage(prompt_tokens=10, completion_tokens=20),
                is_final=True,
            )

        gateway = MagicMock(spec=LLMGateway)
        gateway.chat = AsyncMock(side_effect=[invalid_plan])
        gateway.chat_stream = mock_chat_stream

        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["react"],
        )

        events = []
        async for event in engine.execute_stream(
            messages=[{"role": "user", "content": "Stream task"}],
            tools=[tool],
        ):
            events.append(event)

        event_types = [e.event_type for e in events]
        assert "planning" in event_types
        assert "final_answer" in event_types

    @pytest.mark.asyncio
    async def test_stream_all_fallbacks_exhausted(self):
        """流式模式下所有回退策略耗尽 → RuntimeError"""
        call_count = 0

        class FailingStreamGateway:
            def __init__(self):
                self.chat = AsyncMock(side_effect=self._respond_or_fail)
                self.chat_stream = AsyncMock(side_effect=self._stream_fail)

            async def _respond_or_fail(self, **kwargs):
                nonlocal call_count
                call_count += 1
                if call_count == 1:
                    return make_response(content="Not a plan")
                raise RuntimeError("LLM unavailable")

            async def _stream_fail(self, **kwargs):
                raise RuntimeError("Stream unavailable")

        gateway = FailingStreamGateway()
        engine = ReWOOEngine(llm_gateway=gateway)

        with pytest.raises(RuntimeError, match="All ReWOO fallback strategies exhausted"):
            async for _ in engine.execute_stream(
                messages=[{"role": "user", "content": "Impossible task"}],
            ):
                pass


# ---------------------------------------------------------------------------
# Test 5: Invalid strategy validation
# ---------------------------------------------------------------------------


class TestFallbackStrategyValidation:
    """回退策略参数验证"""

    @pytest.mark.asyncio
    async def test_invalid_strategy_skipped(self):
        """无效策略被跳过，有效策略仍然执行"""
        invalid_plan = make_response(content="Cannot plan")
        direct_response = make_response(content="Direct answer")

        gateway = make_mock_gateway([invalid_plan, direct_response])
        # "invalid_strategy" 应被跳过，"direct" 应执行
        engine = ReWOOEngine(
            llm_gateway=gateway,
            fallback_strategies=["invalid_strategy", "direct"],
        )

        # "invalid_strategy" 被跳过，所以 fallback_strategies 只包含 "direct"
        assert engine._fallback_strategies == ["direct"]

        result = await engine.execute(
            messages=[{"role": "user", "content": "Task"}],
        )

        assert result.output == "Direct answer"
        assert result.fallback_strategy == "direct"

    @pytest.mark.asyncio
    async def test_all_invalid_strategies_uses_defaults(self):
        """全部无效策略时回退到默认策略"""
        engine = ReWOOEngine(
            llm_gateway=MagicMock(spec=LLMGateway),
            fallback_strategies=["invalid1", "invalid2"],
        )

        # 全部无效时回退到默认策略
        assert engine._fallback_strategies == ["simplified_rewoo", "react", "direct"]
