ReMe/reme/steps/index/traverse.py

"""Build a bounded, frontend-ready graph from indexed Markdown wikilinks."""

from collections import deque
from dataclasses import dataclass
from pathlib import Path

from ..base_step import BaseStep
from ...components import R
from ...schema import FileNode, TraverseGraph, TraverseGraphEdge, TraverseGraphNode

_DIRECTION_ALIASES = {
    "out": "forward",
    "forward": "forward",
    "in": "backward",
    "backward": "backward",
    "both": "both",
}


@dataclass(frozen=True, slots=True)
class _TraversalLink:
    """An adjacency entry carrying both traversal and original edge direction."""

    neighbor: str
    source: str
    target: str
    target_anchor: str | None


Adjacency = dict[str, list[_TraversalLink]]
EdgeKey = tuple[str, str, str | None]


def _build_adjacency(nodes: list[FileNode]) -> tuple[Adjacency, Adjacency]:
    """Build forward/backward adjacency while preserving actual wikilink direction."""
    forward: Adjacency = {}
    backward: Adjacency = {}
    for node in nodes:
        for link in node.links:
            if not link.target_path:
                continue
            item = _TraversalLink(
                neighbor=link.target_path,
                source=node.path,
                target=link.target_path,
                target_anchor=link.target_anchor,
            )
            forward.setdefault(node.path, []).append(item)
            backward.setdefault(link.target_path, []).append(
                _TraversalLink(
                    neighbor=node.path,
                    source=item.source,
                    target=item.target,
                    target_anchor=item.target_anchor,
                ),
            )
    return forward, backward


def _traverse(
    seeds: list[str],
    max_depth: int,
    direction: str,
    forward: Adjacency,
    backward: Adjacency,
) -> tuple[dict[str, int], dict[EdgeKey, int]]:
    """Return shortest node depths and directed edges reached by bounded BFS."""
    adjacency_maps = []
    if direction in {"forward", "both"}:
        adjacency_maps.append(forward)
    if direction in {"backward", "both"}:
        adjacency_maps.append(backward)

    node_depths = {seed: 0 for seed in seeds}
    edge_depths: dict[EdgeKey, int] = {}
    queue = deque(seeds)

    while queue:
        current = queue.popleft()
        current_depth = node_depths[current]
        if current_depth >= max_depth:
            continue

        next_depth = current_depth + 1
        for adjacency in adjacency_maps:
            for link in adjacency.get(current, ()):
                edge_key = (link.source, link.target, link.target_anchor)
                previous_edge_depth = edge_depths.get(edge_key)
                if previous_edge_depth is None or next_depth < previous_edge_depth:
                    edge_depths[edge_key] = next_depth

                previous_node_depth = node_depths.get(link.neighbor)
                if previous_node_depth is None or next_depth < previous_node_depth:
                    node_depths[link.neighbor] = next_depth
                    queue.append(link.neighbor)

    return node_depths, edge_depths


def _build_graph(
    *,
    seeds: list[str],
    max_depth: int,
    direction: str,
    indexed_nodes: list[FileNode],
    node_depths: dict[str, int],
    edge_depths: dict[EdgeKey, int],
) -> TraverseGraph:
    """Materialize deterministic public graph schemas from traversal state."""
    node_by_path = {node.path: node for node in indexed_nodes}
    graph_nodes = []
    for path, node_depth in sorted(node_depths.items(), key=lambda item: (item[1], item[0])):
        node = node_by_path.get(path)
        graph_nodes.append(
            TraverseGraphNode(
                id=path,
                path=path,
                name=node.front_matter.name if node is not None else "",
                description=node.front_matter.description if node is not None else "",
                depth=node_depth,
                indexed=node is not None,
            ),
        )

    graph_edges = [
        TraverseGraphEdge(source=source, target=target, target_anchor=anchor, depth=edge_depth)
        for (source, target, anchor), edge_depth in sorted(
            edge_depths.items(),
            key=lambda item: (item[1], item[0][0], item[0][1], item[0][2] or ""),
        )
    ]
    return TraverseGraph(
        seeds=seeds,
        depth=max_depth,
        direction=direction,
        nodes=graph_nodes,
        edges=graph_edges,
    )


@R.register("traverse_step")
class TraverseStep(BaseStep):
    """Return a bounded wikilink graph rooted at one or more workspace paths."""

    async def execute(self):
        assert self.context is not None
        raw_paths = self.context.get("path")
        items = [raw_paths] if isinstance(raw_paths, (str, Path)) else list(raw_paths or [])
        seeds = list(dict.fromkeys(str(path).replace("\\", "/") for path in items if path))
        if not seeds:
            raise ValueError("path is required")

        raw_depth = self.context.get("depth")
        max_depth = 1 if raw_depth is None else int(raw_depth)
        if max_depth < 0:
            raise ValueError("depth must be greater than or equal to 0")

        raw_direction = str(self.context.get("direction") or "both").lower()
        direction = _DIRECTION_ALIASES.get(raw_direction)
        if direction is None:
            raise ValueError(f"direction must be one of {sorted(_DIRECTION_ALIASES)}, got {raw_direction!r}")

        indexed_nodes = await self.file_store.get_nodes()
        forward, backward = _build_adjacency(indexed_nodes)
        node_depths, edge_depths = _traverse(seeds, max_depth, direction, forward, backward)
        graph = _build_graph(
            seeds=seeds,
            max_depth=max_depth,
            direction=direction,
            indexed_nodes=indexed_nodes,
            node_depths=node_depths,
            edge_depths=edge_depths,
        )

        self.context.response.success = True
        self.context.response.answer = graph.model_dump()
        self.logger.info(
            f"[{self.name}] seeds={seeds!r} depth={max_depth} direction={direction} "
            f"nodes={len(graph.nodes)} edges={len(graph.edges)}",
        )
        return self.context.response