"""Wave planner: groups tasks into parallelizable execution waves."""

from __future__ import annotations

from loom.graph.task import Task


def compute_waves(tasks: list[Task], done_ids: set[str] | None = None) -> list[list[Task]]:
    """Topological sort into execution waves.

    Wave 0 = tasks with no pending deps (ready now).
    Wave N = tasks whose deps are all in waves < N or already done.

    Tasks that are already claimed are placed in wave 0 (in-progress).
    """
    done = done_ids or set()
    task_map = {t.id: t for t in tasks}
    resolved: set[str] = set(done)  # IDs whose wave is finalized
    assigned: dict[str, int] = {}

    # Claimed tasks are already in-progress — wave 0
    for t in tasks:
        if t.status == "claimed":
            assigned[t.id] = 0

    wave = 0
    while True:
        batch: list[Task] = []
        for t in tasks:
            if t.id in assigned:
                continue
            all_met = all(
                dep_id in resolved or dep_id not in task_map
                for dep_id in t.depends_on
            )
            if all_met:
                batch.append(t)
        if not batch and not any(tid for tid, w in assigned.items() if tid not in resolved):
            break
        for t in batch:
            assigned[t.id] = wave
        # Mark all tasks assigned up to this wave as resolved for next wave
        resolved.update(tid for tid, w in assigned.items() if w <= wave)
        wave += 1

    # Build wave lists
    max_wave = max(assigned.values(), default=-1)
    waves: list[list[Task]] = [[] for _ in range(max_wave + 1)]
    for t in tasks:
        if t.id in assigned:
            waves[assigned[t.id]].append(t)

    # Tasks with circular or unresolvable deps go in a final wave
    unassigned = [t for t in tasks if t.id not in assigned]
    if unassigned:
        waves.append(unassigned)

    return [w for w in waves if w]


def group_by_files(tasks: list[Task]) -> list[list[Task]]:
    """Group tasks within a wave by file overlap for agent assignment.

    Tasks sharing files should go to the same agent to avoid merge conflicts.
    Uses union-find to cluster tasks with overlapping context.files.
    """
    if not tasks:
        return []

    # Build file -> task indices mapping
    file_to_indices: dict[str, list[int]] = {}
    for i, t in enumerate(tasks):
        files = _get_files(t)
        for f in files:
            file_to_indices.setdefault(f, []).append(i)

    # Union-find
    parent = list(range(len(tasks)))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    # Union tasks that share files
    for indices in file_to_indices.values():
        for j in range(1, len(indices)):
            union(indices[0], indices[j])

    # Group by root
    groups: dict[int, list[Task]] = {}
    for i, t in enumerate(tasks):
        root = find(i)
        groups.setdefault(root, []).append(t)

    return list(groups.values())


def _get_files(task: Task) -> list[str]:
    """Extract files list from task context."""
    if isinstance(task.context, dict):
        files = task.context.get("files", [])
        if isinstance(files, list):
            return files
    return []


def build_orchestration_plan(
    tasks: list[Task], done_ids: set[str] | None = None,
) -> dict:
    """Build a full orchestration plan with waves and agent groups."""
    waves = compute_waves(tasks, done_ids)

    plan_waves = []
    for i, wave_tasks in enumerate(waves):
        groups = group_by_files(wave_tasks)
        wave_info: dict = {
            "wave": i,
            "task_count": len(wave_tasks),
            "agent_groups": [],
        }
        for g, group in enumerate(groups):
            files = set()
            for t in group:
                files.update(_get_files(t))
            wave_info["agent_groups"].append({
                "group": g,
                "task_ids": [t.id for t in group],
                "tasks": [{"id": t.id, "title": t.title, "priority": t.priority,
                           "status": t.status} for t in group],
                "shared_files": sorted(files),
            })
        plan_waves.append(wave_info)

    return {
        "total_waves": len(waves),
        "total_tasks": len(tasks),
        "waves": plan_waves,
    }
