import { DemoBox } from "../../src/components/demo/demo-box";

Large trees can have a significant impact on render and usage performance, and cause slow
interactions. Virtualization is a concept, where the performance hit of a very large list of items
is mitigated by only rendering the items that are currently visible in the viewport.

While this is not trivial to do with nested structures like trees, Headless Tree makes this
easy by flattening the tree structure and providing the tree items as flat list. Virtualization
is not a included feature of Headless Tree, but you can easily pass this flat list to any virtualization
library of your choice, and use that to create a tree that only renders the visible items.

In the sample below, `react-virtual` is used to virtualize the tree and render 100k items while
still being performant in rendering and interaction.

```ts jsx
import type { Meta } from "@storybook/react";
import React, {
  forwardRef,
  useImperativeHandle,
  useRef,
  useState,
} from "react";
import {
  TreeState,
  buildProxiedInstance,
  buildStaticInstance,
  dragAndDropFeature,
  hotkeysCoreFeature,
  keyboardDragAndDropFeature,
  selectionFeature,
  syncDataLoaderFeature,
} from "@headless-tree/core";
import { useTree } from "@headless-tree/react";
import { Virtualizer, useVirtualizer } from "@tanstack/react-virtual";
import cn from "classnames";
import { PropsOfArgtype } from "../argtypes";

const meta = {
  title: "React/Scalability/Basic Virtualization",
  tags: ["homepage"],
  argTypes: {
    itemsPerLevel: {
      type: "number",
    },
    openLevels: {
      type: "number",
    },
    useProxyInstances: {
      type: "boolean",
    },
  },
  args: {
    itemsPerLevel: 10,
    openLevels: 4,
    useProxyInstances: true,
  },
} satisfies Meta;

export default meta;

// story-start
const getInitiallyExpandedItemIds = (
  itemsPerLevel: number,
  openLevels: number,
  prefix = "folder",
): string[] => {
  if (openLevels === 0) {
    return [];
  }

  const expandedItems: string[] = [];

  for (let i = 0; i < itemsPerLevel; i++) {
    expandedItems.push(`${prefix}-${i}`);
  }

  return [
    ...expandedItems,
    ...expandedItems.flatMap((itemId) =>
      getInitiallyExpandedItemIds(itemsPerLevel, openLevels - 1, itemId),
    ),
  ];
};

const Inner = forwardRef<Virtualizer<HTMLDivElement, Element>, any>(
  ({ tree }, ref) => {
    const parentRef = useRef<HTMLDivElement | null>(null);

    const virtualizer = useVirtualizer({
      count: tree.getItems().length,
      getScrollElement: () => parentRef.current,
      estimateSize: () => 27,
    });

    useImperativeHandle(ref, () => virtualizer);

    return (
      <div
        ref={parentRef}
        style={{
          height: `400px`,
          overflow: "auto",
        }}
      >
        <div
          {...tree.getContainerProps()}
          className="tree"
          style={{
            height: `${virtualizer.getTotalSize()}px`,
            width: "100%",
            position: "relative",
          }}
        >
          {virtualizer.getVirtualItems().map((virtualItem) => {
            const item = tree.getItems()[virtualItem.index];
            return (
              <button
                {...item.getProps()}
                key={item.getId()}
                style={{
                  position: "absolute",
                  top: 0,
                  left: 0,
                  width: "100%",
                  transform: `translateY(${virtualItem.start}px)`,
                  paddingLeft: `${item.getItemMeta().level * 20}px`,
                }}
              >
                <div
                  className={cn("treeitem", {
                    focused: item.isFocused(),
                    expanded: item.isExpanded(),
                    selected: item.isSelected(),
                    folder: item.isFolder(),
                    drop: item.isDragTarget(),
                  })}
                >
                  {item.getItemName()}
                </div>
              </button>
            );
          })}
          <div style={tree.getDragLineStyle()} className="dragline" />
        </div>
      </div>
    );
  },
);

export const BasicVirtualization = ({
  itemsPerLevel,
  openLevels,
  useProxyInstances,
}: PropsOfArgtype<typeof meta>) => {
  const virtualizer = useRef<Virtualizer<HTMLDivElement, Element> | null>(null);
  const [state, setState] = useState<Partial<TreeState<string>>>(() => ({
    expandedItems: getInitiallyExpandedItemIds(itemsPerLevel, openLevels),
  }));
  const tree = useTree<string>({
    instanceBuilder: useProxyInstances
      ? buildProxiedInstance
      : buildStaticInstance,
    state,
    setState,
    rootItemId: "folder",
    getItemName: (item) => item.getItemData(),
    isItemFolder: (item) => !item.getItemData().endsWith("item"),
    scrollToItem: (item) => {
      virtualizer.current?.scrollToIndex(item.getItemMeta().index);
    },
    canReorder: true,
    indent: 20,
    dataLoader: {
      getItem: (itemId) => itemId,
      getChildren: (itemId) => {
        const items: string[] = [];
        for (let i = 0; i < itemsPerLevel; i++) {
          items.push(`${itemId}-${i}`);
        }
        return items;
      },
    },
    features: [
      syncDataLoaderFeature,
      selectionFeature,
      hotkeysCoreFeature,
      dragAndDropFeature,
      keyboardDragAndDropFeature,
    ],
  });

  return <Inner tree={tree} ref={virtualizer} />;
};
```

:::warning

You likely will want to use proxified item instances instead of static item instances when
using trees with many items. Read this guide to learn more about [Proxy Item Instances](https://headless-tree.lukasbach.com/llm/recipe/proxy-instances.md).
You can use them by setting the `instanceBuilder` tree config option to `buildProxiedInstance`,
a symbol that you can import from `@headless-tree/core`.

:::

If you need to need a ref to the virtualized DOM items, keep in mind that `treeItem.getProps()` also
returns a ref that needs to be assigned to the DOM element. Also keep in mind that the ref function
of a DOM element is called both on mount and unmount, and calling `treeItem.getProps()` will fail
if the item is already unloaded in the tree. Calling `treeItem.getProps()` outside the ref, like the following,
will work:

```ts jsx
const props = item.getProps();
return (
  <button
    {...props}
    key={virtualItem.key}
    data-index={virtualItem.index}
    ref={(r) => {
      virtualizer.measureElement(r);
      props.ref(r);
    }}
  />
);
```