"""U7: 多步研究任务端到端验证

用 mock LLM Gateway 验证 PlanExecEngine + SharedWorkspace + ReActStepExecutor 的完整闭环。
场景："分析竞品并生成报告" → 搜索→分析→生成
"""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

from agentkit.core.plan_exec_engine import PlanExecEngine
from agentkit.core.shared_workspace import SharedWorkspace
from agentkit.llm.protocol import LLMResponse, TokenUsage


def _make_mock_gateway(responses: list[str]) -> MagicMock:
    """Create a mock LLM Gateway that returns preset responses in order.

    Each response is wrapped in an LLMResponse dataclass to match the real
    LLMGateway.chat() return format.
    """
    gateway = MagicMock()
    call_count = 0

    async def _chat(messages, model="default", **kwargs):
        nonlocal call_count
        if call_count < len(responses):
            resp = responses[call_count]
            call_count += 1
            return LLMResponse(
                content=resp,
                model=model,
                usage=TokenUsage(prompt_tokens=10, completion_tokens=len(resp) // 4),
            )
        return LLMResponse(
            content='{"error": "no more responses"}',
            model=model,
            usage=TokenUsage(prompt_tokens=0, completion_tokens=0),
        )

    gateway.chat = AsyncMock(side_effect=_chat)
    return gateway


class TestMultiStepResearchTask:
    """多步研究任务端到端验证"""

    @pytest.mark.asyncio
    async def test_plan_exec_research_task_with_workspace(self):
        """验证 PlanExecEngine 执行多步研究任务，SharedWorkspace 传递状态"""
        workspace = SharedWorkspace()

        # Mock LLM responses:
        # 1. GoalPlanner: decompose "分析竞品并生成报告" into 3 steps
        # 2-4. Step executor responses for each step
        planner_response = json.dumps({
            "goal": "分析竞品并生成报告",
            "steps": [
                {
                    "step_id": "search",
                    "name": "搜索竞品信息",
                    "description": "搜索飞书和钉钉的竞品信息",
                    "dependencies": [],
                    "required_tools": ["web_search"],
                },
                {
                    "step_id": "analyze",
                    "name": "分析竞品对比",
                    "description": "基于搜索结果分析飞书和钉钉的优缺点",
                    "dependencies": ["search"],
                    "required_tools": [],
                },
                {
                    "step_id": "report",
                    "name": "生成报告",
                    "description": "基于分析结果生成竞品对比报告",
                    "dependencies": ["analyze"],
                    "required_tools": [],
                },
            ],
        })

        step_responses = [
            planner_response,
            "搜索完成：飞书是字节跳动的企业协作平台，钉钉是阿里巴巴的企业通讯工具",
            "分析完成：飞书优势在于文档协作，钉钉优势在于即时通讯和考勤管理",
            "报告：飞书与钉钉竞品对比报告\n1. 文档协作：飞书领先\n2. 即时通讯：钉钉领先\n3. 考勤管理：钉钉领先",
        ]

        gateway = _make_mock_gateway(step_responses)

        engine = PlanExecEngine(
            llm_gateway=gateway,
            workspace=workspace,
            max_replans=0,
        )

        result = await engine.execute(
            messages=[{"role": "user", "content": "分析竞品并生成报告"}],
            tools=[],
            model="default",
            agent_name="research_agent",
        )

        # Verify execution completed
        assert result.status in ("success", "partial")
        assert result.output is not None
        assert len(result.output) > 0

        # Verify workspace has step results
        keys = await workspace.list_keys()
        plan_keys = [k for k in keys if k.startswith("plan:")]
        assert len(plan_keys) > 0, f"Expected workspace keys for plan steps, got: {keys}"

    @pytest.mark.asyncio
    async def test_workspace_dependency_chain(self):
        """验证 SharedWorkspace 在步骤间正确传递依赖结果"""
        workspace = SharedWorkspace()

        # Write a dependency result
        await workspace.write(
            "plan:test_plan:step:search:result",
            {"content": "搜索结果：飞书和钉钉的信息"},
            agent_id="search_agent",
        )

        # Read it back
        data = await workspace.read("plan:test_plan:step:search:result")
        assert data is not None
        assert data["value"]["content"] == "搜索结果：飞书和钉钉的信息"
        assert data["version"] == 1

        # Write a second version
        version = await workspace.write(
            "plan:test_plan:step:search:result",
            {"content": "更新后的搜索结果"},
            agent_id="search_agent",
        )
        assert version == 2

        # Verify version incremented
        data = await workspace.read("plan:test_plan:step:search:result")
        assert data["version"] == 2

    @pytest.mark.asyncio
    async def test_step_event_callback(self):
        """验证 step_event_callback 在步骤完成时被调用"""
        events: list[tuple[str, dict]] = []

        async def on_step_event(event_type: str, data: dict):
            events.append((event_type, data))

        planner_response = json.dumps({
            "goal": "简单任务",
            "steps": [
                {
                    "step_id": "step1",
                    "name": "执行步骤1",
                    "description": "执行一个简单步骤",
                    "dependencies": [],
                    "required_tools": [],
                },
            ],
        })

        gateway = _make_mock_gateway([planner_response, "步骤1完成"])

        engine = PlanExecEngine(
            llm_gateway=gateway,
            step_event_callback=on_step_event,
            max_replans=0,
        )

        await engine.execute(
            messages=[{"role": "user", "content": "简单任务"}],
            tools=[],
        )

        # Verify callback was called
        event_types = [e[0] for e in events]
        assert "plan_generated" in event_types, f"Expected plan_generated event, got: {event_types}"

        # Verify step event contains expected data
        step_events = [e for e in events if e[0] in ("step_completed", "step_failed")]
        assert len(step_events) > 0, f"Expected step events, got: {events}"

    @pytest.mark.asyncio
    async def test_no_workspace_fallback(self):
        """验证 workspace=None 时不影响执行"""
        planner_response = json.dumps({
            "goal": "无workspace任务",
            "steps": [
                {
                    "step_id": "step1",
                    "name": "步骤1",
                    "description": "执行步骤",
                    "dependencies": [],
                    "required_tools": [],
                },
            ],
        })

        gateway = _make_mock_gateway([planner_response, "完成"])

        engine = PlanExecEngine(
            llm_gateway=gateway,
            workspace=None,
            max_replans=0,
        )

        result = await engine.execute(
            messages=[{"role": "user", "content": "无workspace任务"}],
            tools=[],
        )

        assert result.status in ("success", "partial")

    @pytest.mark.asyncio
    async def test_no_callback_fallback(self):
        """验证 callback=None 时不影响执行"""
        planner_response = json.dumps({
            "goal": "无callback任务",
            "steps": [
                {
                    "step_id": "step1",
                    "name": "步骤1",
                    "description": "执行步骤",
                    "dependencies": [],
                    "required_tools": [],
                },
            ],
        })

        gateway = _make_mock_gateway([planner_response, "完成"])

        engine = PlanExecEngine(
            llm_gateway=gateway,
            step_event_callback=None,
            max_replans=0,
        )

        result = await engine.execute(
            messages=[{"role": "user", "content": "无callback任务"}],
            tools=[],
        )

        assert result.status in ("success", "partial")
