import React, { useState } from 'react';
import type { DroppableCollectionReorderEvent, DropPosition } from 'react-aria-components';
import { DropIndicator, useDragAndDrop } from 'react-aria-components';

import type { Complete, Exact, SafeOmit } from '@shoptet/utils';
import { destruct, getControlMode, useStableCallback } from '@shoptet/utils';

import { TreeItem } from './TreeItem/TreeItem';
import type { TreeBaseProps } from './TreeBase';
import { TreeBase } from './TreeBase';
import type { TreeItemKey, TreeNode } from './useTreeState';
import { collectKeys, useTreeState } from './useTreeState';

export type { TreeItemKey, TreeNode } from './useTreeState';
export type { TreeItemRenderProps } from './TreeItem/TreeItem';

export interface TreeMoveEvent<V extends TreeItemKey = string> extends SafeOmit<
  DroppableCollectionReorderEvent,
  'keys' | 'target'
> {
  /** The set of item keys being moved. */
  keys: Set<V>;
  /** The drop target information. */
  target: {
    /** The type of the target (always 'item'). */
    type: 'item';
    /** The key of the target item. */
    key: V;
    /** The position where items should be dropped relative to the target. */
    dropPosition: DropPosition;
  };
}

export type TreeMoveHandler<T extends object = object, V extends TreeItemKey = string> = (
  e: TreeMoveEvent<V>,
  newTree: TreeNode<T, V>[]
) => void | Promise<void>;

export interface TreeOwnProps<T extends object = object, V extends TreeItemKey = string> {
  /** The items to render in the tree (controlled). */
  items?: T[];
  /** The default items to render in the tree (uncontrolled). */
  defaultItems?: T[];
  /**
   * Defines the scope of selection: 'node' selects only the node, 'children' also selects all descendants,
   * 'parents' also selects all ancestors.
   */
  selectionScope?: 'node' | 'children' | 'parents';
  /** Handler called when children need to be loaded for an item. Returns a promise of child items. */
  onLoadChildren?: (key: V) => Promise<T[]>;
  /** Handler called when items are moved via drag and drop. */
  onMove?: TreeMoveHandler<T, V>;
}

export interface InheritedTreeBaseProps<T extends object = object, V extends TreeItemKey = string> extends SafeOmit<
  TreeBaseProps<T, V>,
  'items' | 'onLoadChildren' | 'dragAndDropHooks'
> {}

export interface TreeProps<T extends object = object, V extends TreeItemKey = string>
  extends TreeOwnProps<T, V>, InheritedTreeBaseProps<T, V> {}

export const Tree = Object.assign(
  function Tree<T extends object, V extends TreeItemKey = string>(props: TreeProps<T, V>) {
    const [treeProps, treeBaseProps] = destruct(props, [
      'items',
      'defaultItems',
      'selectionScope',
      'onLoadChildren',
      'onMove',
    ]);

    treeProps satisfies Exact<Complete<TreeOwnProps<T, V>>, typeof treeProps>;
    treeBaseProps satisfies Exact<Complete<InheritedTreeBaseProps<T, V>>, typeof treeBaseProps>;

    const [controlMode] = useState(() => getControlMode(treeProps.items));

    const tree = useTreeState({
      items: treeProps.items,
      defaultItems: treeProps.defaultItems,
      selectedKeys: treeBaseProps.value,
      defaultSelectedKeys: treeBaseProps.defaultValue,
      onSelectedKeysChange: treeBaseProps.onChange,
      getKey: treeBaseProps.getItemValue,
      getChildren: treeBaseProps.getItemChildren,
    });

    const [loadingValues, setLoadingValues] = useState(new Set<V>());
    const [isMoving, setIsMoving] = useState(false);

    const addLoadingValues = useStableCallback((value: V) => {
      setLoadingValues(old => new Set(old).add(value));
    });

    const removeLoadingValue = useStableCallback((value: V) => {
      setLoadingValues(old => {
        const newSet = new Set(old);
        newSet.delete(value);
        return newSet;
      });
    });

    const onLoadChildren = useStableCallback(async (value: V) => {
      addLoadingValues(value);

      try {
        const newChildren = await treeProps.onLoadChildren?.(value);
        const newTree =
          controlMode === 'controlled' ? tree.replace(value, newChildren ?? []) : tree.append(value, newChildren ?? []);
        return collectKeys(newTree.nodeMap.get(value)?.children ?? []);
      } finally {
        removeLoadingValue(value);
      }
    });

    const handleChange = useStableCallback(async (selection: V[]) => {
      if (loadingValues.size > 0) {
        return;
      }

      if (!treeProps.selectionScope || treeProps.selectionScope === 'node') {
        tree.setSelectedKeys?.(selection);
        return;
      }

      const selectionSet = new Set(selection);

      const currentSelection = new Set<V>(tree.selectedKeys ?? []);
      const newSelection = new Set<V>(currentSelection);

      const addedValues = [...selection].filter(value => !currentSelection.has(value));
      const removedValues = [...currentSelection].filter(value => !selectionSet.has(value));

      if (treeProps.selectionScope === 'parents') {
        for (const value of addedValues) {
          newSelection.add(value);

          let parentKey = tree.getItem(value)?.parentKey;

          while (parentKey) {
            newSelection.add(parentKey);

            const parentNode = tree.getItem(parentKey);
            parentKey = parentNode?.parentKey;
          }
        }

        for (const value of removedValues) {
          newSelection.delete(value);
        }
      }

      if (treeProps.selectionScope === 'children') {
        for (const value of addedValues) {
          newSelection.add(value);

          const node = tree.getItem(value);

          if (node && node.children && node.children.length > 0) {
            const childrenValues = collectKeys(node.children);
            childrenValues.forEach(childValue => newSelection.add(childValue));
          } else if (node && onLoadChildren && treeBaseProps.getItemHasChildren?.(node.value)) {
            try {
              const loadedChildrenValues = await onLoadChildren(value);
              loadedChildrenValues.forEach(childValue => newSelection.add(childValue));
            } catch (error) {
              console.debug('Failed to load children for value:', value, error);
            }
          }
        }

        for (const value of removedValues) {
          newSelection.delete(value);

          const node = tree.getItem(value);
          if (node?.children) {
            const childrenValues = collectKeys(node.children);
            childrenValues.forEach(childValue => newSelection.delete(childValue));
          }
        }
      }

      tree.setSelectedKeys?.([...newSelection]);
    });

    const onMove = useStableCallback(async (_e: DroppableCollectionReorderEvent) => {
      const event = _e as TreeMoveEvent<V>;

      if (isMoving) {
        return;
      }

      setIsMoving(true);

      if (event.target.dropPosition === 'before') {
        for (const key of event.keys) {
          const originalPosition = tree.getPosition(key);

          const newTree = tree.moveBefore(event.target.key, [key]);
          try {
            await treeProps.onMove?.(event, newTree?.nodes ?? []);
          } catch {
            if (controlMode === 'controlled') {
              tree.externalSync();
            } else {
              if (originalPosition) {
                tree.move(key, originalPosition.parentKey, originalPosition.index);
              }
            }
          }
        }
      } else if (event.target.dropPosition === 'after') {
        for (const key of event.keys) {
          const originalPosition = tree.getPosition(key);

          const newTree = tree.moveAfter(event.target.key, [key]);

          try {
            await treeProps.onMove?.(event, newTree?.nodes ?? []);
          } catch {
            if (controlMode === 'controlled') {
              tree.externalSync();
            } else {
              if (originalPosition) {
                tree.move(key, originalPosition.parentKey, originalPosition.index);
              }
            }
          }
        }
      } else if (event.target.dropPosition === 'on') {
        for (const key of event.keys) {
          const originalPosition = tree.getPosition(key);

          const targetNode = tree.getItem(event.target.key);

          if (targetNode) {
            const newTree = tree.move(key, event.target.key, targetNode.children ? targetNode.children.length : 0);

            const shouldLoadChildren =
              treeProps.onLoadChildren &&
              (!targetNode.children || targetNode.children.length === 0) &&
              treeBaseProps.getItemHasChildren?.(targetNode.value);

            if (shouldLoadChildren) {
              addLoadingValues(event.target.key);
            }

            try {
              await treeProps.onMove?.(event, newTree?.nodes ?? []);

              if (shouldLoadChildren) {
                try {
                  await onLoadChildren(event.target.key);
                } catch {
                  if (controlMode === 'controlled') {
                    tree.externalSync();
                  } else {
                    tree.remove(key);
                  }
                }
              }
            } catch {
              if (shouldLoadChildren) {
                removeLoadingValue(event.target.key);
              }

              if (controlMode === 'controlled') {
                tree.externalSync();
              } else {
                if (originalPosition) {
                  tree.move(key, originalPosition.parentKey, originalPosition.index);
                }
              }
            }
          }
        }
      }

      setIsMoving(false);
    });

    const { dragAndDropHooks } = useDragAndDrop<TreeNode<T, V>>({
      getItems: treeProps.onMove
        ? (_, items) => [
            {
              'text/plain': items.map(item => treeBaseProps.getItemTextLabel?.(item.value)).join(', '),
            },
          ]
        : undefined,
      onMove,
      renderDropIndicator: target => <DropIndicator target={target} className='v2Tree__dropIndicator' />,
    });

    return (
      <TreeBase
        {...treeBaseProps}
        value={treeBaseProps.value || treeBaseProps.defaultValue ? tree.selectedKeys : undefined}
        onChange={handleChange}
        items={tree.nodes}
        renderItem={
          treeBaseProps.renderItem
            ? props =>
                treeBaseProps.renderItem?.({
                  ...props,
                  item: props.item.value,
                })
            : undefined
        }
        getItemValue={item => item.key}
        getItemChildren={item => item.children ?? []}
        getItemTextLabel={item => treeBaseProps.getItemTextLabel(item.value)}
        getItemTextDescription={
          treeBaseProps.getItemTextDescription ? item => treeBaseProps.getItemTextDescription!(item.value) : undefined
        }
        getItemHasChildren={
          treeBaseProps.getItemHasChildren ? item => treeBaseProps.getItemHasChildren!(item.value) : undefined
        }
        onLoadChildren={treeProps.onLoadChildren ? onLoadChildren : undefined}
        loadingValues={loadingValues}
        dragAndDropHooks={dragAndDropHooks}
      />
    );
  },
  {
    Item: TreeItem,
  }
);