"""Tests for the workflow coordinator."""

from __future__ import annotations

from typing import Any

import pytest

from a2a_mesh.coordinator import WorkflowCoordinator
from a2a_mesh.exceptions import CyclicDependencyError, TaskExecutionError
from a2a_mesh.models import (
    ConsensusConfig,
    ConsensusThreshold,
    FanInStrategy,
    Task,
    TaskStatus,
    Workflow,
)


async def mock_executor(task: Task) -> Any:
    """Mock task executor that echoes the input."""
    return f"result:{task.input}"


async def failing_executor(task: Task) -> Any:
    """Mock executor that always fails."""
    raise RuntimeError("Agent is down")


class TestWorkflowCoordinator:
    """Tests for WorkflowCoordinator."""

    @pytest.mark.asyncio
    async def test_simple_linear_workflow(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="linear",
            tasks=[
                Task(name="step1", input="data"),
                Task(name="step2", depends_on=["step1"]),
                Task(name="step3", depends_on=["step2"]),
            ],
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        assert "step1" in result.task_results
        assert "step2" in result.task_results
        assert "step3" in result.task_results

    @pytest.mark.asyncio
    async def test_parallel_tasks(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="parallel",
            tasks=[
                Task(name="a", input="x"),
                Task(name="b", input="y"),
                Task(name="c", input="z"),
                Task(name="merge", depends_on=["a", "b", "c"]),
            ],
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        assert len(result.task_results) == 4

    @pytest.mark.asyncio
    async def test_diamond_dependency(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="diamond",
            tasks=[
                Task(name="top", input="start"),
                Task(name="left", depends_on=["top"]),
                Task(name="right", depends_on=["top"]),
                Task(name="bottom", depends_on=["left", "right"]),
            ],
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        assert len(result.task_results) == 4

    @pytest.mark.asyncio
    async def test_cyclic_dependency_raises(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="cyclic",
            tasks=[
                Task(name="a", depends_on=["c"]),
                Task(name="b", depends_on=["a"]),
                Task(name="c", depends_on=["b"]),
            ],
        )
        with pytest.raises(CyclicDependencyError):
            await coordinator.execute(workflow)

    @pytest.mark.asyncio
    async def test_failed_task_stops_workflow(self) -> None:
        coordinator = WorkflowCoordinator(executor=failing_executor)
        workflow = Workflow(
            name="failing",
            tasks=[
                Task(name="step1", input="data"),
                Task(name="step2", depends_on=["step1"]),
            ],
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.FAILED
        assert "step1" in result.errors

    @pytest.mark.asyncio
    async def test_fan_out(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="fan-out",
            tasks=[Task(name="research", input="topic")],
            fan_out={"research": 3},
            fan_in={"research": FanInStrategy.MERGE},
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        # Fan-out with merge produces a list of results
        assert isinstance(result.task_results["research"], list)
        assert len(result.task_results["research"]) == 3

    @pytest.mark.asyncio
    async def test_fan_in_first(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="fan-in-first",
            tasks=[Task(name="research", input="topic")],
            fan_out={"research": 3},
            fan_in={"research": FanInStrategy.FIRST},
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        # FIRST strategy returns just the first result
        assert result.task_results["research"] == "result:topic"

    @pytest.mark.asyncio
    async def test_consensus_all_agree(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="consensus",
            tasks=[Task(name="review", input="content")],
            consensus={
                "review": ConsensusConfig(
                    agents=3,
                    threshold=ConsensusThreshold.ALL_AGREE,
                ),
            },
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        assert result.task_results["review"] == "result:content"

    @pytest.mark.asyncio
    async def test_consensus_any(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="consensus-any",
            tasks=[Task(name="review", input="x")],
            consensus={
                "review": ConsensusConfig(
                    agents=2,
                    threshold=ConsensusThreshold.ANY,
                ),
            },
        )
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED

    @pytest.mark.asyncio
    async def test_dependency_injection(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(
            name="injection",
            tasks=[
                Task(name="fetch", input="data"),
                Task(name="process", depends_on=["fetch"]),
            ],
        )
        result = await coordinator.execute(workflow)
        # process should receive fetch's result as input
        assert "result:data" in str(result.task_results["process"])

    @pytest.mark.asyncio
    async def test_empty_workflow(self) -> None:
        coordinator = WorkflowCoordinator(executor=mock_executor)
        workflow = Workflow(name="empty", tasks=[])
        result = await coordinator.execute(workflow)
        assert result.status == TaskStatus.COMPLETED
        assert result.task_results == {}
