"""U4: TeamOrchestrator 无依赖子任务并行模式单元测试 (R1-R5).

Covers:
- TeamPlan.get_independent_subtasks() introspection
- _rebalance_independent_subtasks() MAX_EXPERTS enforcement (R4)
- topological_sort 同层并行 (existing behavior, validated here for U4 contract)
- SharedWorkspace 路径唯一性 (phase_id is UUID → no collision)
- 综合阶段等待所有并行子任务完成 (layer-sequential contract)
"""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

from agentkit.core.handoff_transport import InProcessHandoffTransport
from agentkit.core.protocol import TaskResult, TaskStatus
from agentkit.experts.config import ExpertConfig
from agentkit.experts.expert import Expert
from agentkit.experts.orchestrator import TeamOrchestrator
from agentkit.experts.plan import PlanPhase, TeamPlan
from agentkit.experts.team import ExpertTeam

from tests.unit.experts._helpers import (
    make_chat_stream_mock,
    make_execute_stream_mock,
)


# ── Helpers ────────────────────────────────────────────────────────────


def _make_expert_config(name: str, is_lead: bool = False) -> ExpertConfig:
    return ExpertConfig(
        name=name,
        agent_type="expert",
        persona="测试专家",
        thinking_style="逻辑推理",
        bound_skills=["skill_a"],
        is_lead=is_lead,
        task_mode="llm_generate",
        prompt={"identity": "测试"},
        llm={"model": "default"},
    )


def _make_mock_expert(name: str, is_lead: bool = False) -> MagicMock:
    config = _make_expert_config(name=name, is_lead=is_lead)
    expert = MagicMock(spec=Expert)
    expert.config = config
    expert.is_active = True
    expert.team_id = None
    expert.get_capabilities_summary.return_value = {
        "name": name,
        "persona": config.persona,
        "thinking_style": config.thinking_style,
        "bound_skills": config.bound_skills,
        "is_lead": is_lead,
    }
    mock_agent = MagicMock()
    mock_agent.execute = AsyncMock(
        return_value=TaskResult(
            task_id="test",
            agent_name=name,
            status=TaskStatus.COMPLETED.value,
            output_data={"content": f"Result from {name}"},
            error_message=None,
            started_at=None,
            completed_at=None,
        )
    )
    mock_agent.execute_stream = make_execute_stream_mock(f"Result from {name}")
    mock_agent._llm_gateway = None
    expert.agent = mock_agent
    return expert


def _make_team_with_experts(
    expert_names: list[str] | None = None,
    lead_name: str = "lead",
) -> ExpertTeam:
    team = ExpertTeam()
    transport = AsyncMock(spec=InProcessHandoffTransport)
    team._handoff_transport = transport
    if expert_names is None:
        expert_names = [lead_name, "member1", "member2"]
    for name in expert_names:
        is_lead = name == lead_name
        expert = _make_mock_expert(name=name, is_lead=is_lead)
        team._experts[name] = expert
        if is_lead:
            team._lead_expert_name = name
    return team


def _make_phase(
    id: str, name: str, assigned_expert: str, depends_on: list[str] | None = None
) -> PlanPhase:
    return PlanPhase(
        id=id,
        name=name,
        assigned_expert=assigned_expert,
        task_description=f"task for {name}",
        depends_on=depends_on or [],
    )


def _make_mock_llm_gateway_with_phases(
    phases_responses: list[list[dict]],
    synthesis_content: str = "综合结果",
) -> MagicMock:
    """Gateway that returns successive decompositions then synthesis.

    phases_responses: each element is a list of phase dicts to return on
    successive decomposition calls. After exhausted, returns synthesis.
    """
    gateway = AsyncMock()
    decomp_responses = [MagicMock(content=json.dumps(phases)) for phases in phases_responses]
    synth_response = MagicMock(content=synthesis_content)
    call_count = [0]

    async def chat_side_effect(messages, model=None, **kwargs):
        call_count[0] += 1
        if call_count[0] <= len(decomp_responses):
            return decomp_responses[call_count[0] - 1]
        return synth_response

    gateway.chat = AsyncMock(side_effect=chat_side_effect)
    gateway.chat_stream = make_chat_stream_mock(synthesis_content)
    return gateway


# ── TeamPlan.get_independent_subtasks ─────────────────────────────────


class TestGetIndependentSubtasks:
    """U4/R1: TeamPlan.get_independent_subtasks() returns phases with no
    depends_on."""

    def test_returns_empty_for_no_phases(self):
        plan = TeamPlan(task="test", lead_expert="lead")
        assert plan.get_independent_subtasks() == []

    def test_returns_only_independent_phases(self):
        plan = TeamPlan(
            task="test",
            lead_expert="lead",
            phases=[
                _make_phase("p1", "A", "lead", depends_on=[]),
                _make_phase("p2", "B", "member1", depends_on=["p1"]),
                _make_phase("p3", "C", "member2", depends_on=[]),
                _make_phase("p4", "D", "member1", depends_on=["p1", "p3"]),
            ],
        )
        independent = plan.get_independent_subtasks()
        assert len(independent) == 2
        assert {ph.id for ph in independent} == {"p1", "p3"}

    def test_returns_all_when_no_dependencies(self):
        plan = TeamPlan(
            task="test",
            lead_expert="lead",
            phases=[
                _make_phase("p1", "A", "lead", depends_on=[]),
                _make_phase("p2", "B", "member1", depends_on=[]),
                _make_phase("p3", "C", "member2", depends_on=[]),
            ],
        )
        independent = plan.get_independent_subtasks()
        assert len(independent) == 3


# ── topological_sort 同层并行 contract (U4/R1, R3) ────────────────────


class TestParallelLayerContract:
    """U4/R1, R3: independent subtasks land in layer 0 (parallel); synthesis
    waits for all layers to complete (layer-sequential)."""

    def test_three_independent_subtasks_in_layer_0(self):
        plan = TeamPlan(
            task="test",
            lead_expert="lead",
            phases=[
                _make_phase("p1", "A", "lead", depends_on=[]),
                _make_phase("p2", "B", "member1", depends_on=[]),
                _make_phase("p3", "C", "member2", depends_on=[]),
            ],
        )
        layers = plan.topological_sort()
        assert len(layers) == 1
        assert len(layers[0]) == 3
        assert {ph.id for ph in layers[0]} == {"p1", "p2", "p3"}

    def test_shared_dependency_degrades_to_layered(self):
        """R3: subtasks sharing an upstream dependency are NOT parallel with
        that upstream — they land in later layers."""
        plan = TeamPlan(
            task="test",
            lead_expert="lead",
            phases=[
                _make_phase("p1", "Base", "lead", depends_on=[]),
                _make_phase("p2", "A", "member1", depends_on=["p1"]),
                _make_phase("p3", "B", "member2", depends_on=["p1"]),
            ],
        )
        layers = plan.topological_sort()
        assert len(layers) == 2
        assert len(layers[0]) == 1  # only Base
        assert layers[0][0].id == "p1"
        assert len(layers[1]) == 2  # A and B parallel (same layer)
        assert {ph.id for ph in layers[1]} == {"p2", "p3"}

    def test_synthesis_waits_for_all_layers(self):
        """R3: dependent phase in last layer — synthesis cannot run until
        its layer completes (validated via topological_sort layer count)."""
        plan = TeamPlan(
            task="test",
            lead_expert="lead",
            phases=[
                _make_phase("p1", "A", "lead", depends_on=[]),
                _make_phase("p2", "B", "member1", depends_on=["p1"]),
                _make_phase("p3", "C", "member2", depends_on=["p2"]),
            ],
        )
        layers = plan.topological_sort()
        assert len(layers) == 3  # strict sequential — no parallelism


# ── SharedWorkspace 路径唯一性 (U4/R2) ─────────────────────────────────


class TestSharedWorkspacePathUniqueness:
    """U4/R2: parallel subtask output paths must not collide. PlanPhase.id
    is UUID by default → unique per phase, ensuring
    ``{plan_id}/phase/{phase_id}/output`` never overlaps."""

    def test_default_phase_ids_are_unique(self):
        phases = [
            PlanPhase(name=f"phase_{i}", assigned_expert="lead", task_description="t")
            for i in range(5)
        ]
        ids = [ph.id for ph in phases]
        assert len(set(ids)) == len(ids), "phase_ids must be unique"

    def test_independent_subtasks_have_unique_output_paths(self):
        plan = TeamPlan(
            task="test",
            lead_expert="lead",
            phases=[
                PlanPhase(name="A", assigned_expert="lead", task_description="t"),
                PlanPhase(name="B", assigned_expert="member1", task_description="t"),
                PlanPhase(name="C", assigned_expert="member2", task_description="t"),
            ],
        )
        independent = plan.get_independent_subtasks()
        paths = {f"{plan.id}/phase/{ph.id}/output" for ph in independent}
        assert len(paths) == len(independent)


# ── _rebalance_independent_subtasks (U4/R4 MAX_EXPERTS) ────────────────


class TestRebalanceIndependentSubtasks:
    """U4/R4: when Lead over-decomposes independent subtasks beyond
    MAX_INDEPENDENT_SUBTASKS, request re-decompose with merge hint."""

    @pytest.mark.asyncio
    async def test_no_rebalance_when_within_limit(self):
        """5 independent subtasks ≤ 10 → no re-decompose."""
        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)
        # 5 independent phases
        phases = [_make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(5)]
        original_phases = list(phases)

        result = await orchestrator._rebalance_independent_subtasks(
            lead=team.lead_expert, task="test", phases=phases
        )
        assert result is original_phases or result == original_phases

    @pytest.mark.asyncio
    async def test_rebalance_triggers_when_over_limit(self):
        """11 independent subtasks > 10 → re-decompose with merge hint."""
        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)

        # Original: 11 independent (over limit)
        original_phases = [
            _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11)
        ]
        # Re-decomposed: 8 independent (within limit) — merge succeeded
        gateway = _make_mock_llm_gateway_with_phases(
            phases_responses=[
                [  # first call (original _decompose_task already happened —
                    # but _rebalance calls _decompose_task again, returning this)
                    {
                        "name": f"merged_{i}",
                        "assigned_expert": "lead",
                        "task_description": "t",
                        "depends_on": [],
                    }
                    for i in range(8)
                ],
            ]
        )
        team._experts["lead"].agent._llm_gateway = gateway

        result = await orchestrator._rebalance_independent_subtasks(
            lead=team.lead_expert, task="test", phases=original_phases
        )
        # Should return the re-decomposed phases (8 independent)
        assert len(result) == 8
        assert all(not ph.depends_on for ph in result)
        # Original phases should NOT be returned
        assert {ph.id for ph in result}.isdisjoint({ph.id for ph in original_phases})

    @pytest.mark.asyncio
    async def test_rebalance_no_gateway_returns_original(self):
        """No LLM gateway → can't re-decompose, return original."""
        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)
        # gateway is None by default in _make_mock_expert
        original_phases = [
            _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11)
        ]

        result = await orchestrator._rebalance_independent_subtasks(
            lead=team.lead_expert, task="test", phases=original_phases
        )
        assert result is original_phases

    @pytest.mark.asyncio
    async def test_rebalance_still_over_limit_returns_new_anyway(self):
        """Re-decompose still produces >MAX independent → return new phases
        (MAX_PHASES truncation handles it in execute())."""
        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)

        original_phases = [
            _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11)
        ]
        # Re-decompose returns 12 independent (still over)
        gateway = _make_mock_llm_gateway_with_phases(
            phases_responses=[
                [
                    {
                        "name": f"q{i}",
                        "assigned_expert": "lead",
                        "task_description": "t",
                        "depends_on": [],
                    }
                    for i in range(12)
                ],
            ]
        )
        team._experts["lead"].agent._llm_gateway = gateway

        result = await orchestrator._rebalance_independent_subtasks(
            lead=team.lead_expert, task="test", phases=original_phases
        )
        # Should return the 12-phase re-decomposition (truncation happens later)
        assert len(result) == 12

    @pytest.mark.asyncio
    async def test_rebalance_decompose_failure_returns_original(self):
        """Re-decompose raises LLMProviderError → return original phases."""
        from agentkit.core.exceptions import LLMProviderError

        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)

        original_phases = [
            _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11)
        ]
        gateway = AsyncMock()
        gateway.chat = AsyncMock(side_effect=LLMProviderError("LLM down"))
        gateway.chat_stream = make_chat_stream_mock("err")
        team._experts["lead"].agent._llm_gateway = gateway

        result = await orchestrator._rebalance_independent_subtasks(
            lead=team.lead_expert, task="test", phases=original_phases
        )
        assert result is original_phases

    @pytest.mark.asyncio
    async def test_rebalance_empty_redecompose_returns_original(self):
        """Re-decompose returns empty list → return original phases."""
        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)

        original_phases = [
            _make_phase(f"p{i}", f"phase_{i}", "lead", depends_on=[]) for i in range(11)
        ]
        # Gateway returns malformed JSON → _parse_phases returns []
        gateway = AsyncMock()
        bad_response = MagicMock(content="not json")
        gateway.chat = AsyncMock(return_value=bad_response)
        gateway.chat_stream = make_chat_stream_mock("err")
        team._experts["lead"].agent._llm_gateway = gateway

        result = await orchestrator._rebalance_independent_subtasks(
            lead=team.lead_expert, task="test", phases=original_phases
        )
        assert result is original_phases


# ── Integration: full execute() with parallel subtasks ────────────────


class TestExecuteParallelSubtasks:
    """U4/R1, R3, R5: full execute() path with independent subtasks running
    in parallel + synthesis waits for all."""

    @pytest.mark.asyncio
    async def test_three_parallel_subtasks_complete_then_synthesis(self):
        """3 independent subtasks → all run in layer 0 → synthesis runs
        after all complete."""
        team = _make_team_with_experts()
        orchestrator = TeamOrchestrator(team)

        gateway = _make_mock_llm_gateway_with_phases(
            phases_responses=[
                [
                    {
                        "name": "A",
                        "assigned_expert": "lead",
                        "task_description": "a",
                        "depends_on": [],
                    },
                    {
                        "name": "B",
                        "assigned_expert": "member1",
                        "task_description": "b",
                        "depends_on": [],
                    },
                    {
                        "name": "C",
                        "assigned_expert": "member2",
                        "task_description": "c",
                        "depends_on": [],
                    },
                ],
            ]
        )
        team._experts["lead"].agent._llm_gateway = gateway

        result = await orchestrator.execute("并行任务")

        assert result["status"] == "completed"
        plan: TeamPlan = result["plan"]
        layers = plan.topological_sort()
        assert len(layers) == 1  # all 3 in layer 0 (parallel)
        assert len(layers[0]) == 3

        # All phases completed
        completed = plan.completed_phases
        assert len(completed) == 3

        # phase_results has all 3
        phase_results: dict = result["phase_results"]
        assert len(phase_results) == 3