import React, { useMemo, useState } from 'react';
import type { Meta, StoryObj } from '@storybook/react-vite';

import { Button } from '../Button/Button';
import { ContextMenu } from '../ContextMenu/ContextMenu';
import { Icon } from '../Icon/Icon';
import { IconButton } from '../IconButton/IconButton';
import { Text } from '../Typography/Text/Text';

import type { FlatTreeItem, LazyTreeItem, NestedTreeItem } from './fixtures';
import {
  flatTreeItems,
  lazyTreeItems,
  loadLazyTreeItemChildren,
  nestedTreeItems,
  nestedTreeItemsLarge,
} from './fixtures';
import type { TreeMoveHandler } from './Tree';
import { Tree } from './Tree';
import type { TreeItemKey, TreeNode } from './useTreeState';

const meta: Meta<typeof Tree> = {
  title: 'UI kit/Tree',
  component: Tree,
  tags: ['autodocs'],
  argTypes: {},
};

export default meta;
type Story<T extends object, V extends TreeItemKey> = StoryObj<typeof Tree<T, V>>;

const defaultArgs: Story<any, number>['args'] = {
  label: 'File System',
};

/**
 * Basic tree component with nested items. No selection or drag & drop enabled.
 * This is the simplest usage of the Tree component.
 */
export const Basic: Story<NestedTreeItem, number> = {
  args: {
    ...defaultArgs,
    items: nestedTreeItems,
    getItemValue: item => item.id,
    getItemChildren: item => item.children,
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
  },
  parameters: {
    docs: {
      description: {
        story: 'Basic tree with nested items. No selection or drag & drop.',
      },
    },
  },
};

/**
 * Tree with deeply nested items (30 levels). Tests the component's ability to handle
 * extreme nesting depth. Each level contains a single item nested within the previous one.
 */
export const DeeplyNested: Story<NestedTreeItem, number> = {
  render: args => {
    const [treeItems, setTreeItems] = useState<NestedTreeItem[]>(
      (() => {
        const createDeepItem = (level: number, maxLevel: number): NestedTreeItem => {
          if (level > maxLevel) {
            return {
              id: level,
              title: `Level ${level} - Final Item`,
              type: 'file',
              children: [],
            };
          }
          return {
            id: level,
            title: `Level ${level} - Directory`,
            type: 'directory',
            children: [createDeepItem(level + 1, maxLevel)],
          };
        };
        return [createDeepItem(1, 30)];
      })()
    );

    const removeItem = (items: NestedTreeItem[], idToRemove: number): NestedTreeItem[] => {
      return items.reduce<NestedTreeItem[]>((acc, item) => {
        if (item.id === idToRemove) {
          return acc;
        }
        if (item.children && item.children.length > 0) {
          return [...acc, { ...item, children: removeItem(item.children, idToRemove) }];
        }
        return [...acc, item];
      }, []);
    };

    return (
      <Tree
        label={args.label}
        items={treeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        renderItem={({ item, ...rest }) => (
          <Tree.Item
            {...rest}
            actions={
              <>
                <IconButton
                  icon='shp-edit-icon'
                  variant='secondary'
                  size='sm'
                  ghost
                  onClick={() => {
                    console.log(`Edit item: ${item.title}`);
                  }}
                />
                <IconButton
                  icon='shp-delete'
                  variant='destructive'
                  size='sm'
                  ghost
                  onClick={() => {
                    setTreeItems(prev => removeItem(prev, item.id));
                  }}
                />
                <IconButton
                  icon='shp-plus'
                  variant='primary'
                  size='sm'
                  ghost
                  onClick={() => {
                    console.log(`Create a new sub-item: ${item.title}`);
                  }}
                />
              </>
            }
          />
        )}
      />
    );
  },
  args: {
    ...defaultArgs,
  },
  parameters: {
    docs: {
      description: {
        story: 'Tree with 30 levels of nesting to test handling of extreme depth.',
      },
    },
  },
};

/**
 * Tree with uncontrolled selection state. The tree manages its own selection internally.
 * Use `defaultValue` to set initial selection and `onChange` to listen for changes.
 */
export const SelectionUncontrolled: Story<NestedTreeItem, number> = {
  args: {
    ...defaultArgs,
    items: nestedTreeItems,
    getItemValue: item => item.id,
    getItemChildren: item => item.children,
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
    defaultValue: [],
    onChange: selection => {
      console.log('Selected items:', selection);
    },
  },
};

/**
 * Tree with controlled selection state. The parent component controls which items are selected.
 * Use `value` and `onChange` to fully control the selection state.
 */
export const SelectionControlled: Story<NestedTreeItem, number> = {
  render: args => {
    const [selectedKeys, setSelectedKeys] = useState<number[]>([1, 3, 13]);

    return (
      <div>
        <div style={{ marginBottom: '1rem' }}>
          <strong>Selected items:</strong> {selectedKeys.join(', ')}
        </div>
        <Tree
          {...args}
          items={nestedTreeItems}
          getItemValue={item => item.id}
          getItemChildren={item => item.children}
          getItemTextLabel={item => item.title}
          getItemTextDescription={item => item.description}
          value={selectedKeys}
          onChange={keys => {
            setSelectedKeys(keys);
            console.log('Selection changed:', keys);
          }}
        />
      </div>
    );
  },
  args: {
    ...defaultArgs,
  },
};

/**
 * Tree with 'children' selection scope. When you select a node, all its descendants are automatically selected.
 * This is useful for hierarchical data where selecting a parent implies selecting all children.
 */
export const SelectionScopeChildren: Story<NestedTreeItem, number> = {
  args: {
    ...defaultArgs,
    items: nestedTreeItems,
    getItemValue: item => item.id,
    getItemChildren: item => item.children,
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
    defaultValue: [],
    selectionScope: 'children',
    onChange: selection => {
      console.log('Selected items:', selection);
    },
  },
};

/**
 * Tree with 'parents' selection scope. When you select a node, all its ancestors are automatically selected.
 * This is useful when child selection implies parent selection.
 */
export const SelectionScopeParents: Story<NestedTreeItem, number> = {
  args: {
    ...defaultArgs,
    items: nestedTreeItems,
    getItemValue: item => item.id,
    getItemChildren: item => item.children,
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
    defaultValue: [],
    selectionScope: 'parents',
    onChange: selection => {
      console.log('Selected items:', selection);
    },
  },
};

/**
 * Tree with uncontrolled expanded state. The tree manages which nodes are expanded internally.
 * Use `defaultExpandedValues` to set initial expanded nodes and `onExpandedChange` to listen for changes.
 */
export const ExpansionUncontrolled: Story<NestedTreeItem, number> = {
  args: {
    ...defaultArgs,
    items: nestedTreeItems,
    getItemValue: item => item.id,
    getItemChildren: item => item.children,
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
    defaultExpandedValues: [1, 13],
    onExpandedChange: expanded => {
      console.log('Expanded items:', expanded);
    },
  },
};

/**
 * Tree with controlled expanded state. The parent component controls which nodes are expanded.
 * Use `expandedValues` and `onExpandedChange` to fully control the expansion state.
 * This example shows how to implement "Expand All" and "Collapse All" functionality.
 */
export const ExpansionControlled: Story<NestedTreeItem, number> = {
  render: args => {
    const [expandedKeys, setExpandedKeys] = useState<number[]>([1, 2, 13, 14]);

    return (
      <div>
        <div style={{ marginBottom: '1rem', display: 'flex', gap: '0.5rem', alignItems: 'center' }}>
          <strong>Expanded items:</strong> {expandedKeys.join(', ')}
          <Button onClick={() => setExpandedKeys([1, 2, 5, 9, 13, 14, 15, 19, 24, 27])}>Expand All Directories</Button>
          <Button onClick={() => setExpandedKeys([])}>Collapse All</Button>
        </div>
        <Tree
          {...args}
          items={nestedTreeItems}
          getItemValue={item => item.id}
          getItemChildren={item => item.children}
          getItemTextLabel={item => item.title}
          getItemTextDescription={item => item.description}
          expandedValues={expandedKeys}
          onExpandedChange={keys => {
            setExpandedKeys(keys);
            console.log('Expanded changed:', keys);
          }}
        />
      </div>
    );
  },
  args: {
    ...defaultArgs,
  },
};

/**
 * Tree with custom action buttons for each item. Use the `renderItem` prop to customize
 * the appearance of tree items and add custom actions like edit, delete, or create.
 */
export const CustomActions: Story<NestedTreeItem, number> = {
  render: args => {
    const [treeItems, setTreeItems] = useState<NestedTreeItem[]>(nestedTreeItems);

    const removeItem = (items: NestedTreeItem[], idToRemove: number): NestedTreeItem[] => {
      return items.reduce<NestedTreeItem[]>((acc, item) => {
        if (item.id === idToRemove) {
          return acc;
        }
        if (item.children && item.children.length > 0) {
          return [...acc, { ...item, children: removeItem(item.children, idToRemove) }];
        }
        return [...acc, item];
      }, []);
    };

    return (
      <Tree
        label={args.label}
        items={treeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        renderItem={({ item, ...renderProps }) => (
          <Tree.Item
            {...renderProps}
            actions={
              renderProps.screenSize === 'small' ? (
                <ContextMenu>
                  <ContextMenu.Item
                    onAction={() => {
                      console.log(`Edit item: ${item.title}`);
                    }}
                  >
                    <ContextMenu.Item.Cell grow={false}>
                      <Icon size='sm' icon='shp-edit-icon' />
                    </ContextMenu.Item.Cell>
                    <ContextMenu.Item.Cell>
                      <Text>Edit</Text>
                    </ContextMenu.Item.Cell>
                  </ContextMenu.Item>
                  <ContextMenu.Item
                    onAction={() => {
                      setTreeItems(prev => removeItem(prev, item.id));
                    }}
                  >
                    <ContextMenu.Item.Cell grow={false}>
                      <Icon size='sm' icon='shp-delete' />
                    </ContextMenu.Item.Cell>
                    <ContextMenu.Item.Cell>
                      <Text>Delete</Text>
                    </ContextMenu.Item.Cell>
                  </ContextMenu.Item>
                  <ContextMenu.Item
                    onAction={() => {
                      console.log(`Create a new sub-item: ${item.title}`);
                    }}
                  >
                    <ContextMenu.Item.Cell grow={false}>
                      <Icon size='sm' icon='shp-plus' />
                    </ContextMenu.Item.Cell>
                    <ContextMenu.Item.Cell>
                      <Text>Add subcategory</Text>
                    </ContextMenu.Item.Cell>
                  </ContextMenu.Item>
                </ContextMenu>
              ) : (
                <>
                  <IconButton
                    icon='shp-edit-icon'
                    variant='secondary'
                    size='sm'
                    ghost
                    onClick={() => {
                      console.log(`Edit item: ${item.title}`);
                    }}
                  />
                  <IconButton
                    icon='shp-delete'
                    variant='destructive'
                    size='sm'
                    ghost
                    onClick={() => {
                      setTreeItems(prev => removeItem(prev, item.id));
                    }}
                  />
                  <IconButton
                    icon='shp-plus'
                    variant='primary'
                    size='sm'
                    ghost
                    onClick={() => {
                      console.log(`Create a new sub-item: ${item.title}`);
                    }}
                  />
                </>
              )
            }
          />
        )}
      />
    );
  },
  args: {
    ...defaultArgs,
  },
};

/**
 * Tree with a large dataset (2500 items). This demonstrates the component's performance
 * with many nested items. The tree does not use virtualization.
 */
export const CustomActionsLargeDataset: Story<NestedTreeItem, number> = {
  render: args => {
    const [treeItems, setTreeItems] = useState<NestedTreeItem[]>(nestedTreeItemsLarge);

    const removeItem = (items: NestedTreeItem[], idToRemove: number): NestedTreeItem[] => {
      return items.reduce<NestedTreeItem[]>((acc, item) => {
        if (item.id === idToRemove) {
          return acc;
        }
        if (item.children && item.children.length > 0) {
          return [...acc, { ...item, children: removeItem(item.children, idToRemove) }];
        }
        return [...acc, item];
      }, []);
    };

    return (
      <Tree
        label={args.label}
        items={treeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        renderItem={({ item, ...restProps }) => (
          <Tree.Item
            {...restProps}
            actions={
              restProps.screenSize === 'small' ? (
                <ContextMenu>
                  <ContextMenu.Item
                    onAction={() => {
                      console.log(`Edit item: ${item.title}`);
                    }}
                  >
                    <ContextMenu.Item.Cell grow={false}>
                      <Icon size='sm' icon='shp-edit-icon' />
                    </ContextMenu.Item.Cell>
                    <ContextMenu.Item.Cell>
                      <Text>Edit</Text>
                    </ContextMenu.Item.Cell>
                  </ContextMenu.Item>
                  <ContextMenu.Item
                    onAction={() => {
                      setTreeItems(prev => removeItem(prev, item.id));
                    }}
                  >
                    <ContextMenu.Item.Cell grow={false}>
                      <Icon size='sm' icon='shp-delete' />
                    </ContextMenu.Item.Cell>
                    <ContextMenu.Item.Cell>
                      <Text>Delete</Text>
                    </ContextMenu.Item.Cell>
                  </ContextMenu.Item>
                  <ContextMenu.Item
                    onAction={() => {
                      console.log(`Create a new sub-item: ${item.title}`);
                    }}
                  >
                    <ContextMenu.Item.Cell grow={false}>
                      <Icon size='sm' icon='shp-plus' />
                    </ContextMenu.Item.Cell>
                    <ContextMenu.Item.Cell>
                      <Text>Add subcategory</Text>
                    </ContextMenu.Item.Cell>
                  </ContextMenu.Item>
                </ContextMenu>
              ) : (
                <>
                  <IconButton
                    icon='shp-edit-icon'
                    variant='secondary'
                    size='sm'
                    ghost
                    onClick={() => {
                      console.log(`Edit item: ${item.title}`);
                    }}
                  />
                  <IconButton
                    icon='shp-delete'
                    variant='destructive'
                    size='sm'
                    ghost
                    onClick={() => {
                      setTreeItems(prev => removeItem(prev, item.id));
                    }}
                  />
                  <IconButton
                    icon='shp-plus'
                    variant='primary'
                    size='sm'
                    ghost
                    onClick={() => {
                      console.log(`Create a new sub-item: ${item.title}`);
                    }}
                  />
                </>
              )
            }
          />
        )}
      />
    );
  },
  args: {
    ...defaultArgs,
  },
};

/**
 * Tree with lazy loading in uncontrolled mode. Children are loaded on demand when a node is expanded.
 * Use `onLoadChildren` to load data asynchronously when needed.
 */
export const LazyLoadingUncontrolled: Story<LazyTreeItem, number> = {
  args: {
    ...defaultArgs,
  },
  render: args => {
    return (
      <Tree
        label={args.label}
        defaultItems={lazyTreeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children ?? []}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        getItemHasChildren={item => item.hasChildren ?? false}
        onLoadChildren={loadLazyTreeItemChildren}
      />
    );
  },
};

/**
 * Tree that loads entire subtrees at once. When a node is expanded, all its descendants
 * are loaded recursively in a single operation. Useful when you want to load complete
 * branches instead of one level at a time.
 */
export const LazyLoadingSubtree: Story<LazyTreeItem, number> = {
  args: {
    ...defaultArgs,
  },
  render: args => {
    const loadEntireSubtree = async (key: number): Promise<LazyTreeItem[]> => {
      const children = await loadLazyTreeItemChildren(key);

      const loadedChildren = await Promise.all(
        children.map(async child => {
          if (child.hasChildren) {
            const nestedChildren = await loadEntireSubtree(child.id);
            return {
              ...child,
              children: nestedChildren,
            };
          }
          return child;
        })
      );

      return loadedChildren;
    };

    return (
      <Tree
        label={args.label}
        defaultItems={lazyTreeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children ?? []}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        getItemHasChildren={item => item.hasChildren ?? false}
        onLoadChildren={loadEntireSubtree}
      />
    );
  },
};

/**
 * Tree with lazy loading in controlled mode. The parent component manages the tree state
 * and manually updates it when children are loaded. This gives you full control over the
 * loading process and tree structure.
 */
export const LazyLoadingControlled: Story<LazyTreeItem, number> = {
  args: {
    ...defaultArgs,
  },
  render: args => {
    const [treeItems, setTreeItems] = useState(lazyTreeItems);

    const updateTreeWithChildren = (
      items: LazyTreeItem[],
      parentId: number,
      children: LazyTreeItem[]
    ): LazyTreeItem[] => {
      return items.map(item => {
        if (item.id === parentId) {
          return {
            ...item,
            children: children,
          };
        } else if (item.children && item.children.length > 0) {
          return {
            ...item,
            children: updateTreeWithChildren(item.children, parentId, children),
          };
        }
        return item;
      });
    };

    return (
      <Tree
        label={args.label}
        items={treeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children ?? []}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        getItemHasChildren={item => item.hasChildren ?? false}
        onLoadChildren={async (key: TreeItemKey) => {
          const newChildren = await loadLazyTreeItemChildren(key as number);
          setTreeItems(oldItems => updateTreeWithChildren(oldItems, key as number, newChildren));
          return newChildren;
        }}
      />
    );
  },
};

/**
 * Combines lazy loading with 'children' selection scope. When selecting a node,
 * all its descendants are loaded (if needed) and selected automatically.
 */
export const LazyLoadingWithSelection: Story<LazyTreeItem, number> = {
  args: {
    ...defaultArgs,
  },
  render: args => {
    const loadEntireSubtree = async (key: number): Promise<LazyTreeItem[]> => {
      const children = await loadLazyTreeItemChildren(key);

      const loadedChildren = await Promise.all(
        children.map(async child => {
          if (child.hasChildren) {
            const nestedChildren = await loadEntireSubtree(child.id);
            return {
              ...child,
              children: nestedChildren,
            };
          }
          return child;
        })
      );

      return loadedChildren;
    };

    return (
      <Tree
        label={args.label}
        defaultItems={lazyTreeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children ?? []}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        getItemHasChildren={item => item.hasChildren ?? false}
        onLoadChildren={loadEntireSubtree}
        defaultValue={[]}
        selectionScope='children'
        onChange={selection => {
          console.log('Selected items:', selection);
        }}
      />
    );
  },
};

/**
 * Tree with drag & drop in uncontrolled mode. Users can reorder items by dragging them.
 * The tree manages its own state and notifies you of changes via `onMove`.
 */
export const DragAndDropUncontrolled: Story<NestedTreeItem, number> = {
  args: {
    ...defaultArgs,
    defaultItems: nestedTreeItems,
    getItemValue: item => item.id,
    getItemChildren: item => item.children,
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
    onMove: reorderedItems => {
      console.log('Items reordered:', reorderedItems);
    },
  },
};

/**
 * Tree with drag & drop in controlled mode. The parent component receives the new tree structure
 * when items are reordered and can update its state accordingly. This gives you full control
 * over the drag & drop behavior and allows validation before accepting changes.
 */
export const DragAndDropControlled: Story<NestedTreeItem, number> = {
  render: args => {
    const [treeItems, setTreeItems] = useState<NestedTreeItem[]>(nestedTreeItems);
    const [dragCount, setDragCount] = useState(0);

    const extractItems = (nodes: TreeNode<NestedTreeItem, number>[]): NestedTreeItem[] => {
      return nodes.map(node => {
        const item = node.value;
        if (node.children && node.children.length > 0) {
          return {
            ...item,
            children: extractItems(node.children),
          } as NestedTreeItem;
        }
        return item;
      });
    };

    const onMove: TreeMoveHandler<NestedTreeItem, number> = (_, newTree) => {
      const newItems = extractItems(newTree ?? []);
      setTreeItems(newItems);
      setDragCount(prev => prev + 1);
    };

    return (
      <div>
        <div style={{ marginBottom: '1rem', display: 'flex', gap: '0.5rem', alignItems: 'center' }}>
          <strong>Drag operations:</strong> {dragCount}
          <Button
            onClick={() => {
              setTreeItems(nestedTreeItems);
              setDragCount(0);
            }}
          >
            Reset Order
          </Button>
        </div>

        <Tree
          label={args.label}
          items={treeItems}
          onMove={onMove}
          getItemValue={item => item.id}
          getItemChildren={item => item.children ?? []}
          getItemTextLabel={item => item.title}
          getItemTextDescription={item => item.description}
        />
      </div>
    );
  },
  args: {
    ...defaultArgs,
  },
};

/**
 * Tree with flat data structure and drag & drop. Demonstrates working with denormalized data
 * where items reference their parent via `parentId` instead of nesting children.
 * The tree automatically builds the hierarchy from flat data.
 */
export const DragAndDropFlatData: Story<FlatTreeItem, number> = {
  args: {
    label: 'Flat File System',
    defaultItems: flatTreeItems.filter(item => item.parentId === null),
    getItemValue: item => item.id,
    getItemChildren: item => flatTreeItems.filter(child => child.parentId === item.id),
    getItemTextLabel: item => item.title,
    getItemTextDescription: item => item.description,
    onMove: e => {
      console.log('Items reordered:', e);
    },
  },
};

/**
 * Controlled drag & drop with flat data structure. Converts between the tree's nested structure
 * and your flat data format. This example shows how to maintain a flat data structure while
 * using the tree component's drag & drop functionality.
 */
export const DragAndDropFlatDataControlled: Story<FlatTreeItem, number> = {
  render: args => {
    const [treeItems, setTreeItems] = useState(flatTreeItems);
    const [dragCount, setDragCount] = useState(0);

    const extractItems = (nodes: TreeNode<FlatTreeItem, number>[], parentId: number | null = null): FlatTreeItem[] => {
      const result: FlatTreeItem[] = [];

      for (const node of nodes) {
        const item = { ...node.value, parentId } as FlatTreeItem;
        result.push(item);

        if (node.children && node.children.length > 0) {
          const childItems = extractItems(node.children, node.value.id as number);
          result.push(...childItems);
        }
      }

      return result;
    };

    const onMove: TreeMoveHandler<FlatTreeItem, number> = (_, newTree) => {
      setTreeItems(extractItems(newTree ?? []));
      setDragCount(prev => prev + 1);
    };

    return (
      <div>
        <div style={{ marginBottom: '1rem', display: 'flex', gap: '0.5rem', alignItems: 'center' }}>
          <strong>Drag operations:</strong> {dragCount}
          <Button
            onClick={() => {
              setTreeItems(flatTreeItems);
              setDragCount(0);
            }}
          >
            Reset Order
          </Button>
        </div>

        <Tree
          label={args.label}
          items={useMemo(() => treeItems.filter(item => !item.parentId), [treeItems])}
          onMove={onMove}
          getItemValue={item => item.id}
          getItemChildren={item => treeItems.filter(child => child.parentId === item.id)}
          getItemTextLabel={item => item.title}
          getItemTextDescription={item => item.description}
        />
      </div>
    );
  },
  args: {
    ...defaultArgs,
  },
};

/**
 * Combines lazy loading with drag & drop in uncontrolled mode. Children are loaded on demand
 * and users can reorder items. This is useful for large trees where you don't want to load
 * all data upfront but still need drag & drop functionality.
 */
export const DragAndDropWithLazyLoading: Story<LazyTreeItem, number> = {
  args: {
    ...defaultArgs,
  },
  render: args => {
    return (
      <Tree
        label={args.label}
        defaultItems={lazyTreeItems}
        getItemValue={item => item.id}
        getItemChildren={item => item.children ?? []}
        getItemTextLabel={item => item.title}
        getItemTextDescription={item => item.description}
        getItemHasChildren={item => item.hasChildren ?? false}
        onLoadChildren={loadLazyTreeItemChildren}
        onMove={(e, newItems) => {
          console.log('Items reordered:', e, newItems);
        }}
      />
    );
  },
};

/**
 * Combines lazy loading with drag & drop in controlled mode. Gives you full control over
 * both the loading process and drag & drop operations. This is the most flexible setup
 * for complex applications that need validation and custom behavior.
 */
export const DragAndDropWithLazyLoadingControlled: Story<LazyTreeItem, number> = {
  args: {
    ...defaultArgs,
  },
  render: args => {
    const [treeItems, setTreeItems] = useState<LazyTreeItem[]>(lazyTreeItems);
    const [dragCount, setDragCount] = useState(0);

    const extractItems = (nodes: TreeNode<LazyTreeItem, number>[]): LazyTreeItem[] => {
      return nodes.map(node => {
        const item = node.value;
        if (node.children && node.children.length > 0) {
          return {
            ...item,
            children: extractItems(node.children),
          } as LazyTreeItem;
        }
        return item;
      });
    };

    const updateTreeWithChildren = (
      items: LazyTreeItem[],
      parentId: number,
      children: LazyTreeItem[]
    ): LazyTreeItem[] => {
      return items.map(item => {
        if (item.id === parentId) {
          return {
            ...item,
            children: children,
          };
        } else if (item.children && item.children.length > 0) {
          return {
            ...item,
            children: updateTreeWithChildren(item.children, parentId, children),
          };
        }
        return item;
      });
    };

    return (
      <div>
        <div style={{ marginBottom: '1rem', display: 'flex', gap: '0.5rem', alignItems: 'center' }}>
          <strong>Drag operations:</strong> {dragCount}
          <Button
            onClick={() => {
              setTreeItems(lazyTreeItems);
              setDragCount(0);
            }}
          >
            Reset Order
          </Button>
        </div>

        <Tree
          label={args.label}
          items={treeItems}
          getItemValue={item => item.id}
          getItemChildren={item => item.children ?? []}
          getItemTextLabel={item => item.title}
          getItemTextDescription={item => item.description}
          getItemHasChildren={item => item.hasChildren ?? false}
          onLoadChildren={async (key: TreeItemKey) => {
            const newChildren = await loadLazyTreeItemChildren(key as number);
            setTreeItems(oldItems => updateTreeWithChildren(oldItems, key as number, newChildren));
            return newChildren;
          }}
          onMove={(_, newTree) => {
            const updatedItems = extractItems(newTree ?? []);
            setTreeItems(updatedItems);
            setDragCount(prev => prev + 1);
          }}
        />
      </div>
    );
  },
};