velai/controller.py

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

from velai.dataflow.connection import Connection
from velai.dataflow.enums import DataPortState
from velai.dataflow.graph import DataGraph
from velai.graph_undo import GraphUndoStack
from velai.nodes.base_node import BaseNode

logger = logging.getLogger(__name__)


@dataclass
class GraphController:
    """Handle Vue Flow events and keep the DataGraph plus node state in sync.

    This lives in the app layer so the core dataflow package does not
    depend on your concrete node types.
    """

    graph: DataGraph
    undo_stack: GraphUndoStack | None = None

    def handle_event(self, event: dict[str, Any]) -> None:
        event_type = event.get("type")
        raw_payload = event.get("payload")

        if event_type == "connect":
            payload = raw_payload or {}
            self._on_connect(payload)
        elif event_type == "node_moved":
            payload = raw_payload or {}
            self._on_node_moved(payload)
        elif event_type == "node_resized":
            payload = raw_payload or {}
            self._on_node_resized(payload)
        elif event_type == "node_field_changed":
            payload = raw_payload or {}
            self._on_node_field_changed(payload)
        elif event_type == "node_title_changed":
            payload = raw_payload or {}
            self._on_node_title_changed(payload)
        elif event_type == "edges_delete":
            edges = raw_payload or []
            if isinstance(edges, list):
                self._on_edges_delete(edges)
        elif event_type == "edges_change":
            changes = raw_payload or []
            if isinstance(changes, list):
                self._on_edges_change(changes)
        elif event_type == "nodes_delete":
            nodes = raw_payload or []
            if isinstance(nodes, list):
                self._on_nodes_delete(nodes)
        elif event_type == "nodes_change":
            changes = raw_payload or []
            if isinstance(changes, list):
                self._on_nodes_change(changes)
        # other events (graph_cleared, create_node, etc.) are currently ignored
        # because they do not affect the DataGraph directly here

    def _on_connect(self, payload: dict[str, Any]) -> None:
        source_handle = payload.get("sourceHandle") or ""
        target_handle = payload.get("targetHandle") or ""

        def split(handle: str) -> tuple[str, str]:
            if not handle:
                return "", ""
            if ":" in handle:
                node_id, port = handle.split(":", 1)
                return node_id, port
            return "", handle

        src_node_id, src_port_name = split(source_handle)
        tgt_node_id, tgt_port_name = split(target_handle)

        if not src_node_id:
            src_node_id = payload.get("source") or ""
        if not tgt_node_id:
            tgt_node_id = payload.get("target") or ""

        if not src_node_id or not tgt_node_id:
            return

        start_node = self.graph.nodes.get(src_node_id)
        end_node = self.graph.nodes.get(tgt_node_id)
        if start_node is None or end_node is None:
            return

        start_port = start_node.outputs.get(src_port_name) if start_node.outputs is not None else None
        if start_port is None and start_node.inputs is not None:
            start_port = start_node.inputs.get(src_port_name)

        end_port = end_node.inputs.get(tgt_port_name) if end_node.inputs is not None else None
        if end_port is None and end_node.outputs is not None:
            end_port = end_node.outputs.get(tgt_port_name)

        if start_port is None or end_port is None:
            return

        conn = Connection(
            start_node=start_node,
            start_port=start_port,
            end_node=end_node,
            end_port=end_port,
        )

        try:
            self.graph.add_connection(conn)
        except ValueError:
            # datatype mismatch or capacity problems
            return

        # make the end port "dirty" because topology changed
        end_port.state = DataPortState.DIRTY
        end_port.value = None

    def _update_node_position(self, node_id: str, position: dict[str, Any]) -> None:
        node = self.graph.nodes.get(node_id)
        if node is None:
            return

        pos = position or {}
        x = pos.get("x")
        y = pos.get("y")

        if x is not None:
            try:
                node.x = float(x)
            except (TypeError, ValueError):
                pass
        if y is not None:
            try:
                node.y = float(y)
            except (TypeError, ValueError):
                pass

    def _on_node_moved(self, payload: dict[str, Any]) -> None:
        node_id = payload.get("id")
        if not node_id:
            return
        position = payload.get("position") or {}
        self._update_node_position(node_id, position)

    def _on_node_resized(self, payload: dict[str, Any]) -> None:
        node_id = payload.get("id")

        if not node_id:
            return

        node = self.graph.nodes.get(node_id)
        if node is None:
            return

        width = payload.get("width")
        height = payload.get("height")

        if width is not None:
            node.width = float(width)

        if height is not None:
            node.height = float(height)

    def _on_node_field_changed(self, payload: dict[str, Any]) -> None:
        node_id = payload.get("id")
        field = payload.get("field")
        value = payload.get("value")

        if not node_id or not field:
            return

        node = self.graph.nodes.get(node_id)
        if node is None or not isinstance(node, BaseNode):
            return

        node.set_state({field: value})

    def _on_node_title_changed(self, payload: dict[str, Any]) -> None:
        node_id = payload.get("id")
        title = payload.get("title")

        if not node_id:
            return

        node = self.graph.nodes.get(node_id)
        if node is None or not isinstance(node, BaseNode):
            return

        if title and str(title).strip():
            node.data.custom_title = str(title).strip()
        else:
            node.data.custom_title = None

    def _on_edges_delete(self, edges: list[dict[str, Any]]) -> None:
        """Remove matching connections from the DataGraph when edges are deleted in Vue."""
        if not edges:
            return

        for edge in edges:
            if not isinstance(edge, dict):
                continue

            src_id = edge.get("source")
            tgt_id = edge.get("target")
            src_handle = edge.get("sourceHandle") or ""
            tgt_handle = edge.get("targetHandle") or ""

            src_port = src_handle.split(":", 1)[1] if ":" in src_handle else None
            tgt_port = tgt_handle.split(":", 1)[1] if ":" in tgt_handle else None

            def should_remove(conn: Connection) -> bool:
                if src_id and conn.start_node.node_id != src_id:
                    return False
                if tgt_id and conn.end_node.node_id != tgt_id:
                    return False
                if src_port is not None and conn.start_port.name != src_port:
                    return False
                if tgt_port is not None and conn.end_port.name != tgt_port:
                    return False
                return True

            self.graph.connections = [c for c in self.graph.connections if not should_remove(c)]

    def _on_edges_change(self, changes: list[dict[str, Any]]) -> None:
        """Handle generic edge changes.

        Vue Flow sends EdgeChange objects.
        """
        if not changes:
            return

        edges_to_delete: list[dict[str, Any]] = []

        for change in changes:
            if not isinstance(change, dict):
                continue

            if change.get("type") == "remove":
                # this is quite ugly, the edge should be sent directly or read from the list of edges from the graph
                edge = change
                if isinstance(edge, dict):
                    edges_to_delete.append(edge)

        if edges_to_delete:
            self._on_edges_delete(edges_to_delete)

    def _on_nodes_delete(self, nodes: list[dict[str, Any]]) -> None:
        """Remove nodes and all their connections when Vue deletes them."""
        if not nodes:
            return

        node_ids = {n.get("id") for n in nodes if isinstance(n, dict) and n.get("id")}
        self._delete_nodes(node_ids)

    def _on_nodes_change(self, changes: list[dict[str, Any]]) -> None:
        """Handle generic node changes.

        Currently supports:
        - type == "remove": delete node and related connections
        - type == "position": update node position like 'node_moved'
        Other change types (select, dimensions, etc.) do not affect the DataGraph.
        """
        if not changes:
            return

        node_ids_to_delete: set[str] = set()

        for change in changes:
            if not isinstance(change, dict):
                continue

            ctype = change.get("type")
            node_id = change.get("id")

            if ctype == "remove" and node_id:
                node_ids_to_delete.add(node_id)
            elif ctype == "position" and node_id:
                # Vue Flow usually sends 'position' for logical node coordinates.
                position = change.get("position") or change.get("positionAbsolute") or {}
                self._update_node_position(node_id, position)

        if node_ids_to_delete:
            self._delete_nodes(node_ids_to_delete)

    def _delete_nodes(self, node_ids: set[str]) -> None:
        if not node_ids:
            return

        if self.undo_stack is not None:
            self.undo_stack.record_node_deletion(self.graph, node_ids)

        self.graph.connections = [
            c
            for c in self.graph.connections
            if c.start_node.node_id not in node_ids and c.end_node.node_id not in node_ids
        ]

        for node_id in node_ids:
            self.graph.nodes.pop(node_id, None)