turbostarter/apps/web/src/modules/diagram/components/editor/DiagramCanvas.tsx

"use client";

import { useCallback } from "react";
import {
  ReactFlow,
  ReactFlowProvider,
  Background,
  Controls,
  MiniMap,
  BackgroundVariant,
  Panel,
} from "@xyflow/react";
import type { Node } from "@xyflow/react";

import { useGraphStore } from "../../stores/useGraphStore";
import { useAutoLayout } from "../../hooks/useAutoLayout";
import { bfsPath } from "../../lib/bfs-path";
import {
  BpmnActivityNode,
  BpmnSubprocessNode,
  BpmnStartEventNode,
  BpmnEndEventNode,
  BpmnTimerEventNode,
  BpmnMessageEventNode,
  BpmnGatewayNode,
  BpmnDataObjectNode,
  BpmnAnnotationNode,
  BpmnPoolNode,
  BpmnLaneNode,
  BpmnGroupNode,
  BpmnSequenceEdge,
  BpmnMessageEdge,
  BpmnAssociationEdge,
} from "../../types/bpmn";
import { ErEntityNode } from "../../types/er/ErEntityNode";
import { ErRelationshipEdge } from "../../types/er/ErRelationshipEdge";
import { OrgchartPersonNode } from "../../types/orgchart/OrgchartPersonNode";
import { OrgchartHierarchyEdge } from "../../types/orgchart/OrgchartHierarchyEdge";
import { ArchServiceNode } from "../../types/architecture/ArchServiceNode";
import { ArchDatabaseNode } from "../../types/architecture/ArchDatabaseNode";
import { ArchQueueNode } from "../../types/architecture/ArchQueueNode";
import { ArchLoadBalancerNode } from "../../types/architecture/ArchLoadBalancerNode";
import { ArchExternalNode } from "../../types/architecture/ArchExternalNode";
import { ArchConnectionEdge } from "../../types/architecture/ArchConnectionEdge";
import { SeqParticipantNode } from "../../types/sequence/SeqParticipantNode";
import { SeqFragmentNode } from "../../types/sequence/SeqFragmentNode";
import { SeqSyncEdge } from "../../types/sequence/SeqSyncEdge";
import { SeqAsyncEdge } from "../../types/sequence/SeqAsyncEdge";
import { SeqReturnEdge } from "../../types/sequence/SeqReturnEdge";
import { FlowProcessNode } from "../../types/flowchart/FlowProcessNode";
import { FlowDecisionNode } from "../../types/flowchart/FlowDecisionNode";
import { FlowTerminalNode } from "../../types/flowchart/FlowTerminalNode";
import { FlowIoNode } from "../../types/flowchart/FlowIoNode";
import { FlowSubprocessNode } from "../../types/flowchart/FlowSubprocessNode";
import { FlowEdge } from "../../types/flowchart/FlowEdge";

const nodeTypes = {
  bpmnActivity: BpmnActivityNode,
  bpmnSubprocess: BpmnSubprocessNode,
  bpmnStartEvent: BpmnStartEventNode,
  bpmnEndEvent: BpmnEndEventNode,
  bpmnTimerEvent: BpmnTimerEventNode,
  bpmnMessageEvent: BpmnMessageEventNode,
  bpmnGateway: BpmnGatewayNode,
  bpmnDataObject: BpmnDataObjectNode,
  bpmnAnnotation: BpmnAnnotationNode,
  bpmnPool: BpmnPoolNode,
  bpmnLane: BpmnLaneNode,
  bpmnGroup: BpmnGroupNode,
  erEntity: ErEntityNode,
  orgchartPerson: OrgchartPersonNode,
  archService: ArchServiceNode,
  archDatabase: ArchDatabaseNode,
  archQueue: ArchQueueNode,
  archLoadBalancer: ArchLoadBalancerNode,
  archExternal: ArchExternalNode,
  seqParticipant: SeqParticipantNode,
  seqFragment: SeqFragmentNode,
  flowProcess: FlowProcessNode,
  flowDecision: FlowDecisionNode,
  flowTerminal: FlowTerminalNode,
  flowIo: FlowIoNode,
  flowSubprocess: FlowSubprocessNode,
};

const edgeTypes = {
  bpmnSequence: BpmnSequenceEdge,
  bpmnMessage: BpmnMessageEdge,
  bpmnAssociation: BpmnAssociationEdge,
  erRelationship: ErRelationshipEdge,
  orgchartHierarchy: OrgchartHierarchyEdge,
  archConnection: ArchConnectionEdge,
  seqSync: SeqSyncEdge,
  seqAsync: SeqAsyncEdge,
  seqReturn: SeqReturnEdge,
  flowEdge: FlowEdge,
};

/** Container node types that should not participate in BFS highlighting */
const CONTAINER_TYPES = new Set(["bpmnPool", "bpmnLane", "bpmnGroup", "seqFragment"]);

function MarkerDefs() {
  return (
    <svg style={{ position: "absolute", width: 0, height: 0 }}>
      <defs>
        {/* BPMN markers */}
        <marker
          id="bpmn-arrow-filled"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10 Z"
            fill="var(--edge-default, #666)"
          />
        </marker>
        <marker
          id="bpmn-arrow-open"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10"
            fill="none"
            stroke="var(--edge-default, #666)"
            strokeWidth={1.5}
          />
        </marker>
        {/* E-R markers */}
        <marker
          id="er-arrow"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10 Z"
            fill="var(--diagram-er, #7c3aed)"
          />
        </marker>
        {/* Architecture markers */}
        <marker
          id="arch-arrow"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10 Z"
            fill="var(--diagram-architecture, #71717a)"
          />
        </marker>
        {/* Sequence markers */}
        <marker
          id="seq-arrow-filled"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10 Z"
            fill="var(--diagram-sequence, #f59e0b)"
          />
        </marker>
        <marker
          id="seq-arrow-open"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10"
            fill="none"
            stroke="var(--diagram-sequence, #f59e0b)"
            strokeWidth={1.5}
          />
        </marker>
        {/* Flowchart markers */}
        <marker
          id="flow-arrow"
          viewBox="0 0 10 10"
          refX={10}
          refY={5}
          markerWidth={8}
          markerHeight={8}
          orient="auto-start-reverse"
        >
          <path
            d="M 0 0 L 10 5 L 0 10 Z"
            fill="var(--diagram-flowchart, #e11d48)"
          />
        </marker>
      </defs>
    </svg>
  );
}

function CanvasInner() {
  const nodes = useGraphStore((s) => s.nodes);
  const edges = useGraphStore((s) => s.edges);
  const onNodesChange = useGraphStore((s) => s.onNodesChange);
  const onEdgesChange = useGraphStore((s) => s.onEdgesChange);
  const onViewportChange = useGraphStore((s) => s.onViewportChange);
  const highlightedNodeId = useGraphStore((s) => s.highlightedNodeId);
  const setHighlightedNodeId = useGraphStore((s) => s.setHighlightedNodeId);
  const setNodes = useGraphStore((s) => s.setNodes);
  const setEdges = useGraphStore((s) => s.setEdges);

  const { isLayouting } = useAutoLayout();

  const clearHighlight = useCallback(() => {
    if (!highlightedNodeId) return;
    setHighlightedNodeId(null);
    setNodes(nodes.map((n) => ({ ...n, className: undefined })));
    setEdges(edges.map((e) => ({ ...e, className: undefined })));
  }, [highlightedNodeId, nodes, edges, setHighlightedNodeId, setNodes, setEdges]);

  const handleNodeClick = useCallback(
    (_: React.MouseEvent, node: Node) => {
      // Skip container nodes (pools, lanes, groups)
      if (CONTAINER_TYPES.has(node.type ?? "")) return;

      const store = useGraphStore.getState();

      // Toggle off if clicking the same node
      if (store.highlightedNodeId === node.id) {
        store.setHighlightedNodeId(null);
        store.setNodes(
          store.nodes.map((n) => ({ ...n, className: undefined })),
        );
        store.setEdges(
          store.edges.map((e) => ({ ...e, className: undefined })),
        );
        return;
      }

      // Compute BFS path from clicked node
      const graphEdges = store.edges
        .filter((e) => e.type !== "bpmnGroup")
        .map((e) => ({ from: e.source, to: e.target }));
      const { nodeSet, edgeSet } = bfsPath(node.id, graphEdges);

      // Apply highlight/dim classes
      store.setHighlightedNodeId(node.id);
      store.setNodes(
        store.nodes.map((n) => {
          if (CONTAINER_TYPES.has(n.type ?? "")) {
            return { ...n, className: undefined };
          }
          return {
            ...n,
            className: nodeSet.has(n.id) ? "highlighted" : "dimmed",
          };
        }),
      );
      store.setEdges(
        store.edges.map((e) => ({
          ...e,
          className: edgeSet.has(`${e.source}->${e.target}`)
            ? "highlighted"
            : "dimmed",
        })),
      );
    },
    [],
  );

  return (
    <div className="w-full h-full">
      <MarkerDefs />
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onViewportChange={onViewportChange}
        onNodeClick={handleNodeClick}
        onPaneClick={clearHighlight}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        fitView
        colorMode="system"
        proOptions={{ hideAttribution: true }}
      >
        <Background
          variant={BackgroundVariant.Dots}
          gap={20}
          size={1}
          color="var(--canvas-grid)"
        />
        <Controls showInteractive={false} />
        <MiniMap
          pannable
          zoomable
          style={{ width: 120, height: 80 }}
        />
        {isLayouting && (
          <Panel position="top-center">
            <div className="rounded-md bg-background/80 px-3 py-1.5 text-xs text-muted-foreground backdrop-blur-sm border border-border">
              Computing layout...
            </div>
          </Panel>
        )}
      </ReactFlow>
    </div>
  );
}

export function DiagramCanvas() {
  return (
    <ReactFlowProvider>
      <CanvasInner />
    </ReactFlowProvider>
  );
}