feat(ui): dnd mostly working (WIP)

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/ContainerElementComponent.tsx

@@ -1,12 +1,15 @@
 import { Flex, IconButton, type SystemStyleObject } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { ContainerContext, DepthContext } from 'features/nodes/components/sidePanel/builder/contexts';
+import {
+  ContainerContextProvider,
+  DepthContextProvider,
+  useDepthContext,
+} from 'features/nodes/components/sidePanel/builder/contexts';
 import { DividerElementComponent } from 'features/nodes/components/sidePanel/builder/DividerElementComponent';
 import { FormElementEditModeWrapper } from 'features/nodes/components/sidePanel/builder/FormElementEditModeWrapper';
 import { HeadingElementComponent } from 'features/nodes/components/sidePanel/builder/HeadingElementComponent';
 import { NodeFieldElementComponent } from 'features/nodes/components/sidePanel/builder/NodeFieldElementComponent';
 import { TextElementComponent } from 'features/nodes/components/sidePanel/builder/TextElementComponent';
-import { useMonitorForFormElementDnd } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
 import { formElementAdded, selectWorkflowFormMode, useElement } from 'features/nodes/store/workflowSlice';
 import type { ContainerElement } from 'features/nodes/types/workflow';
 import {
@@ -18,7 +21,7 @@ import {
   isNodeFieldElement,
   isTextElement,
 } from 'features/nodes/types/workflow';
-import { memo, useCallback, useContext } from 'react';
+import { memo, useCallback } from 'react';
 import { PiPlusBold } from 'react-icons/pi';
 import type { Equals } from 'tsafe';
 import { assert } from 'tsafe';
@@ -52,34 +55,33 @@ export const ContainerElementComponent = memo(({ id }: { id: string }) => {
 ContainerElementComponent.displayName = 'ContainerElementComponent';
 
 export const ContainerElementComponentViewMode = memo(({ el }: { el: ContainerElement }) => {
-  const depth = useContext(DepthContext);
+  const depth = useDepthContext();
   const { id, data } = el;
   const { children, direction } = data;
 
   return (
-    <DepthContext.Provider value={depth + 1}>
-      <ContainerContext.Provider value={data}>
+    <DepthContextProvider depth={depth + 1}>
+      <ContainerContextProvider id={id} direction={direction}>
         <Flex id={id} className={CONTAINER_CLASS_NAME} sx={sx} data-container-direction={direction}>
           {children.map((childId) => (
             <FormElementComponent key={childId} id={childId} />
           ))}
         </Flex>
-      </ContainerContext.Provider>{' '}
-    </DepthContext.Provider>
+      </ContainerContextProvider>
+    </DepthContextProvider>
   );
 });
 ContainerElementComponentViewMode.displayName = 'ContainerElementComponentViewMode';
 
 export const ContainerElementComponentEditMode = memo(({ el }: { el: ContainerElement }) => {
-  const depth = useContext(DepthContext);
+  const depth = useDepthContext();
   const { id, data } = el;
   const { children, direction } = data;
-  useMonitorForFormElementDnd(id, children);
 
   return (
     <FormElementEditModeWrapper element={el}>
-      <DepthContext.Provider value={depth + 1}>
-        <ContainerContext.Provider value={data}>
+      <DepthContextProvider depth={depth + 1}>
+        <ContainerContextProvider id={id} direction={direction}>
           <Flex id={id} className={CONTAINER_CLASS_NAME} sx={sx} data-container-direction={direction}>
             {children.map((childId) => (
               <FormElementComponent key={childId} id={childId} />
@@ -87,8 +89,8 @@ export const ContainerElementComponentEditMode = memo(({ el }: { el: ContainerEl
             {direction === 'row' && children.length < 3 && depth < 2 && <AddColumnButton containerId={id} />}
             {direction === 'column' && depth < 1 && <AddRowButton containerId={id} />}
           </Flex>
-        </ContainerContext.Provider>
-      </DepthContext.Provider>
+        </ContainerContextProvider>
+      </DepthContextProvider>
     </FormElementEditModeWrapper>
   );
 });

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/DividerElementComponent.tsx

@@ -1,12 +1,12 @@
 import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { Flex } from '@invoke-ai/ui-library';
 import { useAppSelector } from 'app/store/storeHooks';
-import { ContainerContext } from 'features/nodes/components/sidePanel/builder/contexts';
+import { useContainerContext } from 'features/nodes/components/sidePanel/builder/contexts';
 import { FormElementEditModeWrapper } from 'features/nodes/components/sidePanel/builder/FormElementEditModeWrapper';
 import { selectWorkflowFormMode, useElement } from 'features/nodes/store/workflowSlice';
 import type { DividerElement } from 'features/nodes/types/workflow';
 import { DIVIDER_CLASS_NAME, isDividerElement } from 'features/nodes/types/workflow';
-import { memo, useContext } from 'react';
+import { memo } from 'react';
 
 const sx: SystemStyleObject = {
   bg: 'base.700',
@@ -40,7 +40,7 @@ export const DividerElementComponent = memo(({ id }: { id: string }) => {
 DividerElementComponent.displayName = 'DividerElementComponent';
 
 export const DividerElementComponentViewMode = memo(({ el }: { el: DividerElement }) => {
-  const container = useContext(ContainerContext);
+  const container = useContainerContext();
   const { id } = el;
 
   return (
@@ -56,7 +56,7 @@ export const DividerElementComponentViewMode = memo(({ el }: { el: DividerElemen
 DividerElementComponentViewMode.displayName = 'DividerElementComponentViewMode';
 
 export const DividerElementComponentEditMode = memo(({ el }: { el: DividerElement }) => {
-  const container = useContext(ContainerContext);
+  const container = useContainerContext();
   const { id } = el;
 
   return (

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/DndListDropIndicator.tsx

@@ -0,0 +1,122 @@
+// Adapted from https://github.com/alexreardon/pdnd-react-tailwind/blob/main/src/drop-indicator.tsx
+import type { Edge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/types';
+import type { SystemStyleObject } from '@invoke-ai/ui-library';
+import { Box } from '@invoke-ai/ui-library';
+import type { DndListTargetState } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
+
+/**
+ * Design decisions for the drop indicator's main line
+ */
+const line = {
+  thickness: 2,
+  backgroundColor: 'base.500',
+};
+
+type DropIndicatorProps = {
+  /**
+   * The `edge` to draw a drop indicator on.
+   *
+   * `edge` is required as for the best possible performance
+   * outcome you should only render this component when it needs to do something
+   *
+   * @example {closestEdge && <DropIndicator edge={closestEdge} />}
+   */
+  edge: Edge;
+  /**
+   * `gap` allows you to position the drop indicator further away from the drop target.
+   * `gap` should be the distance between your drop targets
+   * a drop indicator will be rendered halfway between the drop targets
+   * (the drop indicator will be offset by half of the `gap`)
+   *
+   * `gap` should be a valid CSS length.
+   * @example "8px"
+   * @example "var(--gap)"
+   */
+  gap?: string;
+};
+
+const lineStyles: SystemStyleObject = {
+  display: 'block',
+  position: 'absolute',
+  zIndex: 1,
+  borderRadius: 'full',
+  // Blocking pointer events to prevent the line from triggering drag events
+  // Dragging over the line should count as dragging over the element behind it
+  pointerEvents: 'none',
+  background: line.backgroundColor,
+};
+
+type Orientation = 'horizontal' | 'vertical';
+
+const orientationStyles: Record<Orientation, SystemStyleObject> = {
+  horizontal: {
+    height: `${line.thickness}px`,
+    left: 2,
+    right: 2,
+  },
+  vertical: {
+    width: `${line.thickness}px`,
+    top: 2,
+    bottom: 2,
+  },
+};
+
+const edgeToOrientationMap: Record<Edge, Orientation> = {
+  top: 'horizontal',
+  bottom: 'horizontal',
+  left: 'vertical',
+  right: 'vertical',
+};
+
+const edgeStyles: Record<Edge, SystemStyleObject> = {
+  top: {
+    top: 'var(--local-line-offset)',
+  },
+  right: {
+    right: 'var(--local-line-offset)',
+  },
+  bottom: {
+    bottom: 'var(--local-line-offset)',
+  },
+  left: {
+    left: 'var(--local-line-offset)',
+  },
+};
+
+/**
+ * __Drop indicator__
+ *
+ * A drop indicator is used to communicate the intended resting place of the draggable item. The orientation of the drop indicator should always match the direction of the content flow.
+ */
+function DndDropIndicatorInternal({ edge, gap = '0px' }: DropIndicatorProps) {
+  /**
+   * To clearly communicate the resting place of a draggable item during a drag operation,
+   * the drop indicator should be positioned half way between draggable items.
+   */
+  const lineOffset = `calc(-0.5 * (${gap} + ${line.thickness}px))`;
+  const orientation = edgeToOrientationMap[edge];
+
+  return (
+    <Box
+      sx={{ ...lineStyles, ...orientationStyles[orientation], ...edgeStyles[edge], '--local-line-offset': lineOffset }}
+    />
+  );
+}
+
+export const DndListDropIndicator = ({ dndState, gap }: { dndState: DndListTargetState; gap?: string }) => {
+  if (dndState.type !== 'is-dragging-over') {
+    return null;
+  }
+
+  if (!dndState.closestCenterOrEdge || dndState.closestCenterOrEdge === 'center') {
+    return null;
+  }
+
+  return (
+    <DndDropIndicatorInternal
+      edge={dndState.closestCenterOrEdge}
+      // This is the gap between items in the list, used to calculate the position of the drop indicator
+      gap={gap || 'var(--invoke-space-2)'}
+    />
+  );
+};

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/FormElementEditModeWrapper.tsx

@@ -1,18 +1,20 @@
 import { Flex, type FlexProps, IconButton, Spacer, Text } from '@invoke-ai/ui-library';
 import { useAppDispatch } from 'app/store/storeHooks';
 import { getPrefixedId } from 'features/controlLayers/konva/util';
-import { DndListDropIndicator } from 'features/dnd/DndListDropIndicator';
-import type { DndListTargetState } from 'features/dnd/types';
-import { DepthContext } from 'features/nodes/components/sidePanel/builder/contexts';
+import { useContainerContext, useDepthContext } from 'features/nodes/components/sidePanel/builder/contexts';
+import { DndListDropIndicator } from 'features/nodes/components/sidePanel/builder/DndListDropIndicator';
+import type { DndListTargetState } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
 import { useDraggableFormElement } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
 import { formElementRemoved } from 'features/nodes/store/workflowSlice';
 import { type FormElement, isContainerElement } from 'features/nodes/types/workflow';
 import { startCase } from 'lodash-es';
-import { memo, useCallback, useContext, useRef } from 'react';
+import { memo, useCallback, useRef } from 'react';
 import { PiXBold } from 'react-icons/pi';
 
 export const EDIT_MODE_WRAPPER_CLASS_NAME = getPrefixedId('edit-mode-wrapper', '-');
 
+export const getEditModeWrapperId = (id: string) => `${id}-edit-mode-wrapper`;
+
 const getHeaderBgColor = (depth: number) => {
   if (depth <= 1) {
     return 'base.800';
@@ -40,6 +42,9 @@ const getBgColor = (dndListState: DndListTargetState) => {
     case 'is-dragging':
       return 'red';
     case 'is-dragging-over':
+      if (dndListState.closestCenterOrEdge === 'center') {
+        return 'magenta';
+      }
       return 'blue';
     case 'preview':
       return 'green';
@@ -50,22 +55,27 @@ export const FormElementEditModeWrapper = memo(
   ({ element, children, ...rest }: { element: FormElement } & FlexProps) => {
     const draggableRef = useRef<HTMLDivElement>(null);
     const dragHandleRef = useRef<HTMLDivElement>(null);
-    const [dndListState] = useDraggableFormElement(element.id, draggableRef, dragHandleRef);
-    const depth = useContext(DepthContext);
+    const container = useContainerContext();
+    const [dndListState] = useDraggableFormElement(element.id, container?.id ?? null, draggableRef, dragHandleRef);
+    const depth = useDepthContext();
     const dispatch = useAppDispatch();
     const removeElement = useCallback(() => {
       dispatch(formElementRemoved({ id: element.id }));
     }, [dispatch, element.id]);
 
+    if (dndListState.type !== 'idle') {
+      // console.log(element.id, 'dndListState', dndListState);
+    }
+
     return (
       <Flex
+        id={getEditModeWrapperId(element.id)}
         ref={draggableRef}
         position="relative"
         className={EDIT_MODE_WRAPPER_CLASS_NAME}
         flexDir="column"
-        borderWidth={1}
+        boxShadow="0 0 0 1px var(--invoke-colors-base-750)"
         borderRadius="base"
-        borderColor="base.750"
         alignItems="center"
         justifyContent="flex-start"
         w="full"
@@ -81,13 +91,14 @@ export const FormElementEditModeWrapper = memo(
           h={8}
           bg={getHeaderBgColor(depth)}
           borderTopRadius="inherit"
-          borderBottomWidth={1}
+          // borderBottomWidth={1}
           borderColor="inherit"
           alignItems="center"
           cursor="grab"
         >
           <Text fontWeight="semibold" noOfLines={1} wordBreak="break-all">
-            {getHeaderLabel(element)}
+            {element.id}
+            {/* {getHeaderLabel(element)} */}
           </Text>
           <Spacer />
           <IconButton

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/TextElementComponent.tsx

@@ -40,11 +40,11 @@ export const TextElementComponentEditMode = memo(({ el }: { el: TextElement }) =
   const { content, fontSize } = data;
 
   return (
-    <Flex id={id} className={TEXT_CLASS_NAME}>
-      <FormElementEditModeWrapper element={el}>
+    <FormElementEditModeWrapper element={el}>
+      <Flex id={id} className={TEXT_CLASS_NAME}>
         <Text fontSize={fontSize}>{content}</Text>
-      </FormElementEditModeWrapper>
-    </Flex>
+      </Flex>
+    </FormElementEditModeWrapper>
   );
 });
 TextElementComponentEditMode.displayName = 'TextElementComponentEditMode';

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/WorkflowBuilder.tsx

@@ -2,6 +2,7 @@ import { Button, Flex } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
 import { FormElementComponent } from 'features/nodes/components/sidePanel/builder/ContainerElementComponent';
+import { useMonitorForFormElementDnd } from 'features/nodes/components/sidePanel/builder/use-builder-dnd';
 import { formLoaded, formModeToggled, selectWorkflowFormMode } from 'features/nodes/store/workflowSlice';
 import { elements, rootElementId } from 'features/nodes/types/workflow';
 import { memo, useCallback, useEffect } from 'react';
@@ -9,14 +10,16 @@ import { memo, useCallback, useEffect } from 'react';
 export const WorkflowBuilder = memo(() => {
   const dispatch = useAppDispatch();
   const mode = useAppSelector(selectWorkflowFormMode);
+  useMonitorForFormElementDnd();
 
   useEffect(() => {
+    // dispatch(formReset());
     dispatch(formLoaded({ elements, rootElementId }));
   }, [dispatch]);
   return (
     <ScrollableContent>
       <Flex w="full" justifyContent="center">
-        <Flex flexDir="column" w={mode === 'view' ? '768px' : 'min-content'} minW='768px'>
+        <Flex flexDir="column" w={mode === 'view' ? '512px' : 'min-content'} minW="512px">
           <ToggleModeButton />
           {rootElementId && <FormElementComponent id={rootElementId} />}
         </Flex>

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/center-or-closest-edge.ts

@@ -0,0 +1,71 @@
+// Adapted from https://github.com/atlassian/pragmatic-drag-and-drop/blob/main/packages/hitbox/src/closest-edge.ts
+// This adaptation adds 'center' as a possible target
+import type { Input, Position } from '@atlaskit/pragmatic-drag-and-drop/types';
+import type { Edge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/dist/types/types';
+
+export type CenterOrEdge = 'center' | Edge;
+
+// re-exporting type to make it easy to use
+
+const getDistanceToCenterOrEdge: {
+  [TKey in CenterOrEdge]: (rect: DOMRect, client: Position) => number;
+} = {
+  top: (rect, client) => Math.abs(client.y - rect.top),
+  right: (rect, client) => Math.abs(rect.right - client.x),
+  bottom: (rect, client) => Math.abs(rect.bottom - client.y),
+  left: (rect, client) => Math.abs(client.x - rect.left),
+  center: (rect, client) => {
+    const centerX = rect.left + rect.width / 2;
+    const centerY = rect.top + rect.height / 2;
+    return Math.sqrt((client.x - centerX) ** 2 + (client.y - centerY) ** 2);
+  },
+};
+
+// using a symbol so we can guarantee a key with a unique value
+const uniqueKey = Symbol('centerWithClosestEdge');
+
+/**
+ * Adds a unique `Symbol` to the `userData` object. Use with `extractClosestEdge()` for type safe lookups.
+ */
+export function attachClosestCenterOrEdge(
+  userData: Record<string | symbol, unknown>,
+  {
+    element,
+    input,
+    allowedCenterOrEdge,
+  }: {
+    element: Element;
+    input: Input;
+    allowedCenterOrEdge: CenterOrEdge[];
+  }
+): Record<string | symbol, unknown> {
+  const client: Position = {
+    x: input.clientX,
+    y: input.clientY,
+  };
+  // I tried caching the result of `getBoundingClientRect()` for a single
+  // frame in order to improve performance.
+  // However, on measurement I saw no improvement. So no longer caching
+  const rect: DOMRect = element.getBoundingClientRect();
+  const entries = allowedCenterOrEdge.map((edge) => {
+    return {
+      edge,
+      value: getDistanceToCenterOrEdge[edge](rect, client),
+    };
+  });
+
+  // edge can be `null` when `allowedCenterOrEdge` is []
+  const addClosestCenterOrEdge: CenterOrEdge | null = entries.sort((a, b) => a.value - b.value)[0]?.edge ?? null;
+
+  return {
+    ...userData,
+    [uniqueKey]: addClosestCenterOrEdge,
+  };
+}
+
+/**
+ * Returns the value added by `attachClosestEdge()` to the `userData` object. It will return `null` if there is no value.
+ */
+export function extractClosestCenterOrEdge(userData: Record<string | symbol, unknown>): CenterOrEdge | null {
+  return (userData[uniqueKey] as CenterOrEdge) ?? null;
+}

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/contexts.ts

@@ -1,5 +0,0 @@
-import type { ContainerElement } from 'features/nodes/types/workflow';
-import { createContext } from 'react';
-
-export const ContainerContext = createContext<ContainerElement['data'] | null>(null);
-export const DepthContext = createContext<number>(0);

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/contexts.tsx

@@ -0,0 +1,35 @@
+import type { ContainerElement, ElementId } from 'features/nodes/types/workflow';
+import type { PropsWithChildren } from 'react';
+import { createContext, memo, useContext, useMemo } from 'react';
+
+type ContainerContextValue = {
+  id: ElementId;
+  direction: ContainerElement['data']['direction'];
+};
+
+const ContainerContext = createContext<ContainerContextValue | null>(null);
+
+export const ContainerContextProvider = memo(
+  ({ id, direction, children }: PropsWithChildren<ContainerContextValue>) => {
+    const ctxValue = useMemo(() => ({ id, direction }), [id, direction]);
+    return <ContainerContext.Provider value={ctxValue}>{children}</ContainerContext.Provider>;
+  }
+);
+ContainerContextProvider.displayName = 'ContainerContextProvider';
+
+export const useContainerContext = () => {
+  const container = useContext(ContainerContext);
+  return container;
+};
+
+const DepthContext = createContext<number>(0);
+
+export const DepthContextProvider = memo(({ depth, children }: PropsWithChildren<{ depth: number }>) => {
+  return <DepthContext.Provider value={depth}>{children}</DepthContext.Provider>;
+});
+DepthContextProvider.displayName = 'DepthContextProvider';
+
+export const useDepthContext = () => {
+  const depth = useContext(DepthContext);
+  return depth;
+};

invokeai/frontend/web/src/features/nodes/components/sidePanel/builder/use-builder-dnd.ts

@@ -4,105 +4,173 @@ import {
   dropTargetForElements,
   monitorForElements,
 } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
-import { attachClosestEdge, extractClosestEdge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/closest-edge';
-import { reorderWithEdge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/util/reorder-with-edge';
 import { getStore } from 'app/store/nanostores/store';
 import { useAppDispatch } from 'app/store/storeHooks';
 import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
-import type { DndListTargetState } from 'features/dnd/types';
-import { idle } from 'features/dnd/types';
 import { firefoxDndFix, triggerPostMoveFlash } from 'features/dnd/util';
-import { formElementContainerDataChanged } from 'features/nodes/store/workflowSlice';
-import type { ElementId, FormElement } from 'features/nodes/types/workflow';
+import type { CenterOrEdge } from 'features/nodes/components/sidePanel/builder/center-or-closest-edge';
+import {
+  attachClosestCenterOrEdge,
+  extractClosestCenterOrEdge,
+} from 'features/nodes/components/sidePanel/builder/center-or-closest-edge';
+import { getEditModeWrapperId } from 'features/nodes/components/sidePanel/builder/FormElementEditModeWrapper';
+import { formElementMoved } from 'features/nodes/store/workflowSlice';
+import type { ContainerElement, ElementId, FormElement } from 'features/nodes/types/workflow';
 import { isContainerElement } from 'features/nodes/types/workflow';
 import type { RefObject } from 'react';
 import { useEffect, useState } from 'react';
 import { flushSync } from 'react-dom';
 import { assert } from 'tsafe';
 
-export const useMonitorForFormElementDnd = (containerId: string, children: ElementId[]) => {
+/**
+ * States for a dnd list with containers.
+ */
+export type DndListTargetState =
+  | {
+      type: 'idle';
+    }
+  | {
+      type: 'preview';
+      container: HTMLElement;
+    }
+  | {
+      type: 'is-dragging';
+    }
+  | {
+      type: 'is-dragging-over';
+      closestCenterOrEdge: CenterOrEdge | null;
+    };
+export const idle: DndListTargetState = { type: 'idle' };
+
+type DndData = {
+  element: FormElement;
+  container: ContainerElement | null;
+};
+
+const getElement = <T extends FormElement>(id: ElementId, guard?: (el: FormElement) => el is T): T => {
+  const el = getStore().getState().workflow.form?.elements[id];
+  assert(el);
+  if (guard) {
+    assert(guard(el));
+    return el;
+  } else {
+    return el as T;
+  }
+};
+
+const adjustIndexForDrop = (index: number, edge: Exclude<CenterOrEdge, 'center'>) => {
+  if (edge === 'left' || edge === 'top') {
+    return index - 1;
+  }
+  return index + 1;
+};
+
+export const useMonitorForFormElementDnd = () => {
   const dispatch = useAppDispatch();
 
   useEffect(() => {
     return monitorForElements({
-      canMonitor({ source }) {
-        return (source.data as FormElement).id === containerId;
-      },
+      // canMonitor({ source }) {
+      //   return (source.data as FormElement).id === containerId;
+      // },
+      canMonitor: () => true,
       onDrop({ location, source }) {
         const target = location.current.dropTargets[0];
         if (!target) {
           return;
         }
 
-        const sourceData = source.data as FormElement;
-        const targetData = target.data as FormElement;
+        const sourceData = source.data as DndData;
+        const targetData = target.data as DndData;
 
-        const sourceElementId = sourceData.id;
-        const targetElementId = targetData.id;
+        const sourceElementId = sourceData.element.id;
+        const targetElementId = targetData.element.id;
 
-        const childrenClone = [...children];
+        const closestCenterOrEdge = extractClosestCenterOrEdge(targetData);
 
-        const indexOfSource = childrenClone.findIndex((elementId) => elementId === sourceElementId);
-        const indexOfTarget = childrenClone.findIndex((elementId) => elementId === targetElementId);
+        if (closestCenterOrEdge === 'center') {
+          const targetContainer = getElement(targetElementId);
+          if (!isContainerElement(targetContainer)) {
+            // Shouldn't happen - when dropped on the center of drop target, the target should always be a container type.
+            return;
+          }
+          flushSync(() => {
+            dispatch(formElementMoved({ id: sourceElementId, containerId: targetContainer.id }));
+          });
+        } else if (closestCenterOrEdge) {
+          if (targetData.container) {
+            const targetContainer = getElement(targetData.container.id);
+            if (!isContainerElement(targetContainer)) {
+              // Shouldn't happen - drop targets should always have a container.
+              return;
+            }
+            const indexOfSource = targetContainer.data.children.findIndex((elementId) => elementId === sourceElementId);
+            const indexOfTarget = targetContainer.data.children.findIndex((elementId) => elementId === targetElementId);
 
-        if (indexOfTarget < 0 || indexOfSource < 0) {
+            if (indexOfSource === indexOfTarget) {
+              // Don't move if the source and target are the same index, meaning same position in the list.
+              return;
+            }
+
+            const adjustedIndex = adjustIndexForDrop(indexOfTarget, closestCenterOrEdge);
+
+            if (indexOfSource === adjustedIndex) {
+              // Don't move if the source is already in the correct position.
+              return;
+            }
+
+            flushSync(() => {
+              dispatch(
+                formElementMoved({
+                  id: sourceElementId,
+                  containerId: targetContainer.id,
+                  index: indexOfTarget,
+                })
+              );
+            });
+          }
+        } else {
+          // No container, cannot do anything
           return;
         }
+        // const childrenClone = [...targetData.container.data.children];
 
-        // Don't move if the source and target are the same index, meaning same position in the list
-        if (indexOfSource === indexOfTarget) {
-          return;
-        }
+        // const indexOfSource = childrenClone.findIndex((elementId) => elementId === sourceElementId);
+        // const indexOfTarget = childrenClone.findIndex((elementId) => elementId === targetElementId);
 
-        const closestEdgeOfTarget = extractClosestEdge(targetData);
-
-        // It's possible that the indices are different, but refer to the same position. For example, if the source is
-        // at 2 and the target is at 3, but the target edge is 'top', then the entity is already in the correct position.
-        // We should bail if this is the case.
-        // let edgeIndexDelta = 0;
-
-        // if (closestEdgeOfTarget === 'bottom') {
-        //   edgeIndexDelta = 1;
-        // } else if (closestEdgeOfTarget === 'top') {
-        //   edgeIndexDelta = -1;
-        // }
-
-        // If the source is already in the correct position, we don't need to move it.
-        // if (indexOfSource === indexOfTarget + edgeIndexDelta) {
+        // if (indexOfTarget < 0 || indexOfSource < 0) {
         //   return;
         // }
 
-        const reorderedChildren = reorderWithEdge({
-          list: childrenClone,
-          startIndex: indexOfSource,
-          indexOfTarget,
-          closestEdgeOfTarget,
-          axis: 'vertical',
-        });
+        // // Don't move if the source and target are the same index, meaning same position in the list
+        // if (indexOfSource === indexOfTarget) {
+        //   return;
+        // }
 
         // Using `flushSync` so we can query the DOM straight after this line
-        flushSync(() => {
-          dispatch(formElementContainerDataChanged({ id: containerId, changes: { children: reorderedChildren } }));
-        });
+        // flushSync(() => {
+        //   dispatch(
+        //     formElementMoved({
+        //       id: sourceElementId,
+        //       containerId: targetData.container.id,
+        //       index: indexOfTarget,
+        //     })
+        //   );
+        // });
 
         // Flash the element that was moved
-        const element = document.querySelector(`#${sourceElementId}`);
+        const element = document.querySelector(`#${getEditModeWrapperId(sourceElementId)}`);
         if (element instanceof HTMLElement) {
           triggerPostMoveFlash(element, colorTokenToCssVar('base.700'));
         }
       },
     });
-  }, [children, containerId, dispatch]);
-};
-
-const getElement = (id: ElementId) => {
-  const el = getStore().getState().workflow.form?.elements[id];
-  assert(el !== undefined);
-  return el;
+  }, [dispatch]);
 };
 
 export const useDraggableFormElement = (
   elementId: ElementId,
+  containerId: ElementId | null,
   draggableRef: RefObject<HTMLElement>,
   dragHandleRef: RefObject<HTMLElement>
 ) => {
@@ -118,14 +186,16 @@ export const useDraggableFormElement = (
     return combine(
       firefoxDndFix(draggableElement),
       draggable({
+        canDrag: () => Boolean(containerId),
         element: draggableElement,
         dragHandle: dragHandleElement,
         getInitialData() {
-          return getElement(elementId);
+          const data: DndData = {
+            element: getElement(elementId),
+            container: containerId ? getElement(containerId, isContainerElement) : null,
+          };
+          return data;
         },
-        // getInitialData() {
-        //   return singleWorkflowFieldDndSource.getData({ fieldIdentifier });
-        // },
         onDragStart() {
           setListDndState({ type: 'is-dragging' });
           setIsDragging(true);
@@ -137,35 +207,57 @@ export const useDraggableFormElement = (
       }),
       dropTargetForElements({
         element: draggableElement,
-        canDrop() {
-          return isContainerElement(getElement(elementId));
-        },
+        // canDrop() {},
         getData({ input }) {
-          const data = { elementId };
-          return attachClosestEdge(data, {
+          const element = getElement(elementId);
+          const container = containerId ? getElement(containerId, isContainerElement) : null;
+
+          const data: DndData = {
+            element,
+            container,
+          };
+
+          const allowedCenterOrEdge: CenterOrEdge[] = [];
+
+          if (isContainerElement(element)) {
+            allowedCenterOrEdge.push('center');
+          }
+
+          if (container?.data.direction === 'row') {
+            allowedCenterOrEdge.push('left', 'right');
+          }
+
+          if (container?.data.direction === 'column') {
+            allowedCenterOrEdge.push('top', 'bottom');
+          }
+
+          return attachClosestCenterOrEdge(data, {
             element: draggableElement,
             input,
-            allowedEdges: ['top', 'bottom', 'left', 'right'],
+            allowedCenterOrEdge,
           });
         },
         getIsSticky() {
           return true;
         },
-        onDragEnter({ self }) {
-          const closestEdge = extractClosestEdge(self.data);
-          setListDndState({ type: 'is-dragging-over', closestEdge });
-          console.log('onDragEnter', self.data);
-        },
-        onDrag({ self }) {
-          const closestEdge = extractClosestEdge(self.data);
+        onDrag({ self, location }) {
+          const innermostDropTargetElement = location.current.dropTargets.at(0)?.element;
+
+          // If the innermost target is not this draggable element, bail. We only want to react when dragging over _this_ element.
+          if (!innermostDropTargetElement || innermostDropTargetElement !== draggableElement) {
+            setListDndState(idle);
+            return;
+          }
+
+          const closestCenterOrEdge = extractClosestCenterOrEdge(self.data);
 
           // Only need to update react state if nothing has changed.
           // Prevents re-rendering.
           setListDndState((current) => {
-            if (current.type === 'is-dragging-over' && current.closestEdge === closestEdge) {
+            if (current.type === 'is-dragging-over' && current.closestCenterOrEdge === closestCenterOrEdge) {
               return current;
             }
-            return { type: 'is-dragging-over', closestEdge };
+            return { type: 'is-dragging-over', closestCenterOrEdge };
           });
         },
         onDragLeave() {
@@ -176,7 +268,7 @@ export const useDraggableFormElement = (
         },
       })
     );
-  }, [dragHandleRef, draggableRef, elementId]);
+  }, [containerId, dragHandleRef, draggableRef, elementId]);
 
   return [dndListState, isDragging] as const;
 };