from itertools import groupby

from .compiler import CompiledDelete, CompiledPod, CompiledRange, CompiledRangeErrors
from .types import ExistingPod


def _get_level(r: CompiledRange, pod_to_range: dict[int, CompiledRange]) -> int:
    if isinstance(r.source_pod, ExistingPod):
        return 0
    parent = pod_to_range.get(r.source_pod.pod_id)
    return 0 if parent is None else _get_level(parent, pod_to_range) + 1


def solve_create_ordering(
    ranges: list[CompiledRange | CompiledDelete | CompiledRangeErrors],
) -> list[list[CompiledPod]]:
    """Return waves of pod creates. All pods within a wave are independent and can
    run concurrently; each wave must complete before the next begins."""
    create_ranges = [r for r in ranges if isinstance(r, CompiledRange)]

    pod_to_range: dict[int, CompiledRange] = {pod.pod_id: r for r in create_ranges for pod in r.pods}

    def key(r: CompiledRange) -> int:
        return _get_level(r, pod_to_range)

    sorted_ranges = sorted(create_ranges, key=key)

    return [[pod for r in group for pod in r.pods] for _, group in groupby(sorted_ranges, key=key)]
