xml-pipeline/bloxserver/domain/edges.py

"""
Edge model for connections between nodes.

Edges define how messages flow between nodes in a flow.
They can optionally have conditions for conditional routing.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any
from uuid import UUID, uuid4


class EdgeCondition(str, Enum):
    """Types of edge conditions."""

    ALWAYS = "always"  # Always route (default)
    ON_SUCCESS = "on_success"  # Route only if source succeeds
    ON_ERROR = "on_error"  # Route only if source fails
    CONDITIONAL = "conditional"  # Custom condition expression


@dataclass
class Edge:
    """
    Connection between two nodes.

    Edges define the flow of messages between nodes. In organism.yaml terms,
    they define the `peers` list for agents and the routing paths.
    """

    id: UUID
    source_node_id: UUID
    target_node_id: UUID

    # Optional label for UI
    label: str | None = None

    # Condition for when this edge is followed
    condition: EdgeCondition = EdgeCondition.ALWAYS

    # Custom condition expression (when condition == CONDITIONAL)
    # Example: "payload.status == 'approved'"
    condition_expression: str | None = None

    # Visual properties for canvas
    source_handle: str | None = None  # Which output port
    target_handle: str | None = None  # Which input port
    animated: bool = False
    edge_type: str = "default"  # default, smoothstep, step, straight

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization (React Flow format)."""
        return {
            "id": str(self.id),
            "source": str(self.source_node_id),
            "target": str(self.target_node_id),
            "label": self.label,
            "data": {
                "condition": self.condition.value,
                "conditionExpression": self.condition_expression,
            },
            "sourceHandle": self.source_handle,
            "targetHandle": self.target_handle,
            "animated": self.animated,
            "type": self.edge_type,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Edge:
        """Create edge from dictionary (React Flow format)."""
        edge_data = data.get("data", {})

        # Handle condition
        condition_str = edge_data.get("condition", "always")
        try:
            condition = EdgeCondition(condition_str)
        except ValueError:
            condition = EdgeCondition.ALWAYS

        return cls(
            id=UUID(data["id"]) if isinstance(data.get("id"), str) else data.get("id", uuid4()),
            source_node_id=UUID(data["source"]) if isinstance(data.get("source"), str) else data["source"],
            target_node_id=UUID(data["target"]) if isinstance(data.get("target"), str) else data["target"],
            label=data.get("label"),
            condition=condition,
            condition_expression=edge_data.get("conditionExpression") or edge_data.get("condition_expression"),
            source_handle=data.get("sourceHandle") or data.get("source_handle"),
            target_handle=data.get("targetHandle") or data.get("target_handle"),
            animated=data.get("animated", False),
            edge_type=data.get("type", "default"),
        )


def compute_peers(
    node_id: UUID,
    edges: list[Edge],
    node_map: dict[UUID, str],
) -> list[str]:
    """
    Compute the peers list for a node based on outgoing edges.

    Args:
        node_id: The node to compute peers for.
        edges: All edges in the flow.
        node_map: Mapping of node IDs to node names.

    Returns:
        List of peer names (target nodes this node can send to).
    """
    peers: list[str] = []

    for edge in edges:
        if edge.source_node_id == node_id:
            target_name = node_map.get(edge.target_node_id)
            if target_name and target_name not in peers:
                peers.append(target_name)

    return peers


def compute_incoming(
    node_id: UUID,
    edges: list[Edge],
    node_map: dict[UUID, str],
) -> list[str]:
    """
    Compute the list of nodes that can send to this node.

    Args:
        node_id: The target node.
        edges: All edges in the flow.
        node_map: Mapping of node IDs to node names.

    Returns:
        List of source node names.
    """
    incoming: list[str] = []

    for edge in edges:
        if edge.target_node_id == node_id:
            source_name = node_map.get(edge.source_node_id)
            if source_name and source_name not in incoming:
                incoming.append(source_name)

    return incoming


def find_entry_nodes(
    nodes: list[UUID],
    edges: list[Edge],
) -> list[UUID]:
    """
    Find nodes with no incoming edges (entry points).

    These are typically where external triggers connect.
    """
    nodes_with_incoming = {edge.target_node_id for edge in edges}
    return [node_id for node_id in nodes if node_id not in nodes_with_incoming]


def find_exit_nodes(
    nodes: list[UUID],
    edges: list[Edge],
) -> list[UUID]:
    """
    Find nodes with no outgoing edges (exit points).

    These are typically terminal handlers or response nodes.
    """
    nodes_with_outgoing = {edge.source_node_id for edge in edges}
    return [node_id for node_id in nodes if node_id not in nodes_with_outgoing]


def detect_cycles(
    nodes: list[UUID],
    edges: list[Edge],
) -> list[list[UUID]]:
    """
    Detect cycles in the flow graph.

    Returns a list of cycles (each cycle is a list of node IDs).
    Cycles are allowed (agents can self-loop) but should be flagged for review.
    """
    cycles: list[list[UUID]] = []

    # Build adjacency list
    adj: dict[UUID, list[UUID]] = {node_id: [] for node_id in nodes}
    for edge in edges:
        if edge.source_node_id in adj:
            adj[edge.source_node_id].append(edge.target_node_id)

    # DFS for cycle detection
    visited: set[UUID] = set()
    rec_stack: set[UUID] = set()
    path: list[UUID] = []

    def dfs(node: UUID) -> None:
        visited.add(node)
        rec_stack.add(node)
        path.append(node)

        for neighbor in adj.get(node, []):
            if neighbor not in visited:
                dfs(neighbor)
            elif neighbor in rec_stack:
                # Found a cycle
                cycle_start = path.index(neighbor)
                cycles.append(path[cycle_start:] + [neighbor])

        path.pop()
        rec_stack.remove(node)

    for node_id in nodes:
        if node_id not in visited:
            dfs(node_id)

    return cycles