"""Generic async DAG walker for MuseHub.

Mirrors ``muse.core.graph.walk_dag`` but accepts an async adjacency function,
making it suitable for SQLAlchemy AsyncSession queries and other async I/O.

The canonical usage pattern is to write an async adjacency closure over a
session or an in-memory dict, then drive the walk with ``walk_dag_async``:

    async def adjacency(commit_id: str) -> list[str]:
        row = await session.get(MusehubCommit, commit_id)
        return row.parent_ids if row else []

    async for cid in walk_dag_async(head, adjacency, max_nodes=10_000):
        ...
"""
from __future__ import annotations

import collections
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable
from typing import Literal, TypeVar

T = TypeVar("T", bound=object)

AdjacencyFn = Callable[[T], Coroutine[None, None, list[T]]]


async def walk_dag_async(
    starts: "T | list[T] | Iterable[T]",
    adjacency: AdjacencyFn,
    *,
    order: "Literal['bfs', 'dfs']" = "bfs",
    exclude: "set[T] | frozenset[T] | None" = None,
    max_nodes: "int | None" = None,
) -> "AsyncIterator[T]":
    """Async BFS/DFS over any DAG with an async adjacency function.

    Parameters
    ----------
    starts:
        A single node, or a list/iterable of start nodes.  Multi-root BFS
        seeds the walk from all of them simultaneously.
    adjacency:
        Async callable: given a node, returns its immediate neighbours.
        Called once per visited node, before the node is yielded.
    order:
        ``"bfs"`` (breadth-first, default) or ``"dfs"`` (depth-first).
    exclude:
        Nodes treated as already-visited.  They and their subtrees are
        never yielded.  The caller's set is never mutated.
    max_nodes:
        Hard cap on the number of nodes yielded.  ``None`` means no limit.

    Yields
    ------
    T
        Each reachable node at most once, in BFS or DFS visit order.
    """
    if order not in ("bfs", "dfs"):
        raise ValueError(f"order must be 'bfs' or 'dfs', got {order!r}")

    if isinstance(starts, (str, bytes)) or not hasattr(starts, "__iter__"):
        start_list: list[T] = [starts]  # type: ignore[list-item]
    else:
        start_list = list(starts)  # type: ignore[arg-type]

    seen: set[T] = set(exclude) if exclude is not None else set()
    yielded = 0
    queue: collections.deque = collections.deque(start_list)

    while queue:
        if max_nodes is not None and yielded >= max_nodes:
            return

        node = queue.popleft() if order == "bfs" else queue.pop()

        if node in seen:
            continue
        seen.add(node)

        neighbours = await adjacency(node)

        yield node
        yielded += 1

        for neighbour in neighbours:
            if neighbour not in seen:
                queue.append(neighbour)