mcpengine/studio/apps/web/components/canvas/ToolCanvas.tsx

'use client';

import React, { useCallback, useMemo } from 'react';
import {
  ReactFlow,
  Background,
  Controls,
  MiniMap,
  type Node,
  type Edge,
  type OnNodesChange,
  type OnEdgesChange,
  type OnConnect,
  type NodeTypes,
  type EdgeTypes,
  BackgroundVariant,
  ConnectionMode,
  addEdge,
} from '@xyflow/react';
import '@xyflow/react/dist/style.css';

import { ToolNode } from './ToolNode';
import { GroupNode } from './GroupNode';
import { ConnectionEdge } from './ConnectionEdge';
import { CanvasToolbar } from './CanvasToolbar';
import { useCanvasState } from '../../hooks/useCanvasState';

import type { ToolDefinition } from '@mcpengine/ai-pipeline/types';

interface ToolCanvasProps {
  tools: ToolDefinition[];
  onToolSelect: (toolName: string | null) => void;
  onToolsChange: (tools: ToolDefinition[]) => void;
}

const nodeTypes: NodeTypes = {
  tool: ToolNode,
  group: GroupNode,
};

const edgeTypes: EdgeTypes = {
  connection: ConnectionEdge,
};

export function ToolCanvas({ tools, onToolSelect, onToolsChange }: ToolCanvasProps) {
  const {
    nodes,
    edges,
    setNodes,
    setEdges,
    selectedNodeId,
    selectNode,
  } = useCanvasState();

  const onNodesChange: OnNodesChange = useCallback(
    (changes) => {
      setNodes(changes);
    },
    [setNodes]
  );

  const onEdgesChange: OnEdgesChange = useCallback(
    (changes) => {
      setEdges(changes);
    },
    [setEdges]
  );

  const onConnect: OnConnect = useCallback(
    (connection) => {
      useCanvasState.setState((state) => ({
        edges: addEdge(
          { ...connection, type: 'connection', animated: true },
          state.edges
        ),
      }));
    },
    []
  );

  const onNodeClick = useCallback(
    (_: React.MouseEvent, node: Node) => {
      selectNode(node.id);
      onToolSelect(node.id);
    },
    [selectNode, onToolSelect]
  );

  const onPaneClick = useCallback(() => {
    selectNode(null);
    onToolSelect(null);
  }, [selectNode, onToolSelect]);

  const defaultEdgeOptions = useMemo(
    () => ({
      type: 'connection',
      animated: true,
    }),
    []
  );

  return (
    <div className="relative h-full w-full bg-[#0A0F1E]">
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        onNodeClick={onNodeClick}
        onPaneClick={onPaneClick}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        defaultEdgeOptions={defaultEdgeOptions}
        connectionMode={ConnectionMode.Loose}
        fitView
        proOptions={{ hideAttribution: true }}
        className="bg-[#0A0F1E]"
      >
        <Background
          variant={BackgroundVariant.Dots}
          gap={20}
          size={1}
          color="#1E293B"
        />
        <Controls
          showZoom={false}
          showFitView={false}
          showInteractive={false}
          className="hidden"
        />
        <MiniMap
          nodeColor={(node) => {
            if (node.id === selectedNodeId) return '#6366F1';
            return '#374151';
          }}
          maskColor="rgba(10, 15, 30, 0.8)"
          className="!bg-gray-900 !border-gray-700 rounded-lg"
          pannable
          zoomable
        />
      </ReactFlow>
      <CanvasToolbar />
    </div>
  );
}