"""Tests for orchestration wave planner."""

from loom.graph.task import Task
from loom.orchestration.waves import compute_waves, group_by_files, build_orchestration_plan


def _task(id: str, deps: list[str] | None = None, status: str = "pending",
          files: list[str] | None = None, priority: str = "p1") -> Task:
    ctx = {"files": files} if files else {}
    return Task(
        id=id, project_id="00000000-0000-0000-0000-000000000000",
        title=f"Task {id}", status=status, priority=priority,
        depends_on=deps or [], context=ctx,
    )


class TestComputeWaves:
    def test_no_tasks(self):
        assert compute_waves([]) == []

    def test_independent_tasks(self):
        tasks = [_task("a"), _task("b"), _task("c")]
        waves = compute_waves(tasks)
        assert len(waves) == 1
        assert len(waves[0]) == 3

    def test_linear_chain(self):
        tasks = [_task("a"), _task("b", ["a"]), _task("c", ["b"])]
        waves = compute_waves(tasks)
        assert len(waves) == 3
        assert [t.id for t in waves[0]] == ["a"]
        assert [t.id for t in waves[1]] == ["b"]
        assert [t.id for t in waves[2]] == ["c"]

    def test_diamond_deps(self):
        tasks = [_task("a"), _task("b", ["a"]), _task("c", ["a"]), _task("d", ["b", "c"])]
        waves = compute_waves(tasks)
        assert len(waves) == 3
        assert [t.id for t in waves[0]] == ["a"]
        assert {t.id for t in waves[1]} == {"b", "c"}
        assert [t.id for t in waves[2]] == ["d"]

    def test_done_deps_skipped(self):
        tasks = [_task("b", ["a"]), _task("c", ["a"])]
        waves = compute_waves(tasks, done_ids={"a"})
        assert len(waves) == 1
        assert {t.id for t in waves[0]} == {"b", "c"}

    def test_claimed_in_wave_zero(self):
        tasks = [_task("a", status="claimed"), _task("b", ["a"])]
        waves = compute_waves(tasks)
        assert len(waves) == 2
        assert waves[0][0].id == "a"
        assert waves[1][0].id == "b"

    def test_all_independent_single_wave(self):
        tasks = [_task("a", priority="p2"), _task("b", priority="p0"), _task("c", priority="p1")]
        waves = compute_waves(tasks)
        assert len(waves) == 1
        assert {t.id for t in waves[0]} == {"a", "b", "c"}


class TestGroupByFiles:
    def test_no_tasks(self):
        assert group_by_files([]) == []

    def test_no_file_overlap(self):
        tasks = [_task("a", files=["x.py"]), _task("b", files=["y.py"])]
        groups = group_by_files(tasks)
        assert len(groups) == 2

    def test_file_overlap_groups_together(self):
        tasks = [
            _task("a", files=["x.py", "y.py"]),
            _task("b", files=["y.py", "z.py"]),
            _task("c", files=["w.py"]),
        ]
        groups = group_by_files(tasks)
        assert len(groups) == 2
        grouped_ids = [sorted(t.id for t in g) for g in groups]
        assert ["a", "b"] in grouped_ids
        assert ["c"] in grouped_ids

    def test_transitive_overlap(self):
        tasks = [
            _task("a", files=["x.py"]),
            _task("b", files=["x.py", "y.py"]),
            _task("c", files=["y.py"]),
        ]
        groups = group_by_files(tasks)
        assert len(groups) == 1
        assert {t.id for t in groups[0]} == {"a", "b", "c"}

    def test_no_files_separate_groups(self):
        tasks = [_task("a"), _task("b")]
        groups = group_by_files(tasks)
        assert len(groups) == 2


class TestBuildOrchestrationPlan:
    def test_full_plan(self):
        tasks = [
            _task("a", files=["x.py"]),
            _task("b", files=["x.py"]),
            _task("c", ["a"], files=["y.py"]),
            _task("d", ["a"], files=["z.py"]),
        ]
        plan = build_orchestration_plan(tasks)
        assert plan["total_waves"] == 2
        assert plan["total_tasks"] == 4
        # Wave 0: a, b (share x.py -> 1 agent group)
        w0 = plan["waves"][0]
        assert w0["task_count"] == 2
        assert len(w0["agent_groups"]) == 1
        # Wave 1: c, d (different files -> 2 agent groups)
        w1 = plan["waves"][1]
        assert w1["task_count"] == 2
        assert len(w1["agent_groups"]) == 2
