mcpengine/studio/apps/web/hooks/useCanvasState.ts

import { create } from 'zustand';
import {
  type Node,
  type Edge,
  type OnNodesChange,
  type OnEdgesChange,
  applyNodeChanges,
  applyEdgeChanges,
} from '@xyflow/react';

import type { ToolDefinition, ToolGroup } from '@mcpengine/ai-pipeline/types';
import type { ToolNodeData } from '../components/canvas/ToolNode';
import type { GroupNodeData } from '../components/canvas/GroupNode';

interface CanvasState {
  // State
  nodes: Node[];
  edges: Edge[];
  selectedNodeId: string | null;
  inspectorOpen: boolean;

  // Node/Edge change handlers (for React Flow)
  setNodes: OnNodesChange;
  setEdges: OnEdgesChange;

  // Actions
  selectNode: (id: string | null) => void;
  deselectAll: () => void;
  addTool: (tool: ToolDefinition, position?: { x: number; y: number }) => void;
  removeTool: (toolName: string) => void;
  updateTool: (toolName: string, updates: Partial<ToolDefinition>) => void;
  toggleInspector: () => void;

  // Initialization
  initializeFromTools: (tools: ToolDefinition[], groups?: ToolGroup[]) => void;
}

const GRID_COLS = 3;
const NODE_WIDTH = 300;
const NODE_HEIGHT = 160;
const GAP_X = 40;
const GAP_Y = 40;

function toolToNode(
  tool: ToolDefinition,
  index: number,
  parentId?: string
): Node {
  const col = index % GRID_COLS;
  const row = Math.floor(index / GRID_COLS);

  return {
    id: tool.name,
    type: 'tool',
    position: {
      x: col * (NODE_WIDTH + GAP_X) + (parentId ? 16 : 0),
      y: row * (NODE_HEIGHT + GAP_Y) + (parentId ? 48 : 0),
    },
    data: {
      tool,
      selected: false,
      enabled: true,
    } satisfies ToolNodeData,
    ...(parentId ? { parentId, extent: 'parent' as const } : {}),
  };
}

function groupToNode(
  group: ToolGroup,
  groupIndex: number,
  toolStartIndex: number
): Node {
  const rows = Math.ceil(group.tools.length / GRID_COLS);
  const width = GRID_COLS * (NODE_WIDTH + GAP_X) + 32;
  const height = rows * (NODE_HEIGHT + GAP_Y) + 64;

  return {
    id: `group-${group.name}`,
    type: 'group',
    position: {
      x: 0,
      y: groupIndex * (height + 60),
    },
    data: {
      group,
      childCount: group.tools.length,
    } satisfies GroupNodeData,
    style: { width, height },
  };
}

export const useCanvasState = create<CanvasState>((set, get) => ({
  nodes: [],
  edges: [],
  selectedNodeId: null,
  inspectorOpen: false,

  setNodes: (changes) => {
    set({
      nodes: applyNodeChanges(changes, get().nodes),
    });
  },

  setEdges: (changes) => {
    set({
      edges: applyEdgeChanges(changes, get().edges),
    });
  },

  selectNode: (id) => {
    set((state) => ({
      selectedNodeId: id,
      inspectorOpen: id !== null,
      nodes: state.nodes.map((n) => ({
        ...n,
        data: {
          ...n.data,
          selected: n.id === id,
        },
      })),
    }));
  },

  deselectAll: () => {
    set((state) => ({
      selectedNodeId: null,
      inspectorOpen: false,
      nodes: state.nodes.map((n) => ({
        ...n,
        data: { ...n.data, selected: false },
      })),
    }));
  },

  addTool: (tool, position) => {
    const existingCount = get().nodes.filter((n) => n.type === 'tool').length;
    const node = toolToNode(tool, existingCount);
    if (position) {
      node.position = position;
    }
    set((state) => ({
      nodes: [...state.nodes, node],
    }));
  },

  removeTool: (toolName) => {
    set((state) => ({
      nodes: state.nodes.filter((n) => n.id !== toolName),
      edges: state.edges.filter(
        (e) => e.source !== toolName && e.target !== toolName
      ),
      selectedNodeId:
        state.selectedNodeId === toolName ? null : state.selectedNodeId,
      inspectorOpen:
        state.selectedNodeId === toolName ? false : state.inspectorOpen,
    }));
  },

  updateTool: (toolName, updates) => {
    set((state) => ({
      nodes: state.nodes.map((n) => {
        if (n.id !== toolName || n.type !== 'tool') return n;
        const currentData = n.data as ToolNodeData;
        return {
          ...n,
          // If name changed, update the node id
          id: updates.name ?? n.id,
          data: {
            ...currentData,
            tool: { ...currentData.tool, ...updates },
          },
        };
      }),
    }));
  },

  toggleInspector: () => {
    set((state) => ({ inspectorOpen: !state.inspectorOpen }));
  },

  initializeFromTools: (tools, groups) => {
    const nodes: Node[] = [];
    const edges: Edge[] = [];

    if (groups && groups.length > 0) {
      let toolIndex = 0;
      groups.forEach((group, gi) => {
        // Create group node
        nodes.push(groupToNode(group, gi, toolIndex));

        // Create tool nodes as children
        group.tools.forEach((tool, ti) => {
          nodes.push(toolToNode(tool, ti, `group-${group.name}`));
          toolIndex++;
        });
      });
    } else {
      // Flat layout — no groups
      tools.forEach((tool, i) => {
        nodes.push(toolToNode(tool, i));
      });
    }

    set({ nodes, edges, selectedNodeId: null, inspectorOpen: false });
  },
}));