import { useRef, useState } from 'react';

import { useControlledState, useStableCallback } from '@shoptet/utils';

export type TreeItemKey = string | number;

export interface TreeStateOptions<T extends object, V extends TreeItemKey = string> {
  /** Items in the tree (controlled) */
  items?: T[];
  /** Default items in the tree (uncontrolled) */
  defaultItems?: T[];
  /** The currently selected keys in the tree (controlled). */
  selectedKeys?: V[];
  /** The default selected keys in the tree (uncontrolled). */
  defaultSelectedKeys?: V[];
  /** Handler that is called when the selection changes. */
  onSelectedKeysChange?: (selection: V[]) => void;
  /** A function that returns a unique key for an item object. */
  getKey: (item: T) => V;
  /** A function that returns the children for an item object. */
  getChildren: (item: T) => T[];
}

export interface TreeNode<T extends object, V extends TreeItemKey = string> {
  /** A unique key for the tree node. */
  key: V;
  /** The key of the parent node. */
  parentKey?: V | null;
  /** The value object for the tree node. */
  value: T;
  /** Children of the tree node. */
  children: TreeNode<T, V>[] | null;
}

export interface TreeState<T extends object, V extends TreeItemKey = string> {
  /** The root nodes in the tree. */
  nodes: TreeNode<T, V>[];
  /** The keys of the currently selected items in the tree. */
  selectedKeys: V[];
  /** Sets the selected keys. */
  setSelectedKeys(selection: V[]): void;
  /** Gets a node from the tree by key. */
  getItem(key: V): TreeNode<T, V> | undefined;
  /** Gets the position of a node in the tree. */
  getPosition(key: V): { parentKey: V | null; index: number } | undefined;
  /** Appends items as children of a parent node. */
  append(parentKey: V | null, items: T[]): TreeDataState<T, V>;
  /** Replaces the children of a parent node with a new subtree. */
  replace(parentKey: V | null, items: T[]): TreeDataState<T, V>;
  /** Removes an item from the tree by its key. */
  remove(...keys: V[]): TreeDataState<T, V> | undefined;
  /** Moves an item within the tree. */
  move(key: V, toParentKey: V | null, index: number): TreeDataState<T, V>;
  /** Moves one or more items before a given key. */
  moveBefore(key: V, keys: V[]): TreeDataState<T, V>;
  /** Moves one or more items after a given key. */
  moveAfter(key: V, keys: V[]): TreeDataState<T, V>;
  /** Synchronizes the tree to the current items. */
  externalSync(): TreeDataState<T, V>;
}

export interface TreeDataState<T extends object, V extends TreeItemKey = string> {
  nodes: TreeNode<T, V>[];
  nodeMap: Map<V, TreeNode<T, V>>;
}

/**
 * Manages state for an immutable tree data structure, and provides convenience methods to update the data over time.
 */
export function useTreeState<T extends object, V extends TreeItemKey = string>({
  items: itemsProp,
  defaultItems: initialItemsProp,
  selectedKeys: selectedKeysProp,
  defaultSelectedKeys: defaultSelectedKeysProp,
  getKey,
  getChildren,
  onSelectedKeysChange: onChange,
}: TreeStateOptions<T, V>): TreeState<T, V> {
  const [tree, setTreeState] = useState<TreeDataState<T, V>>(() =>
    buildTree(itemsProp ?? initialItemsProp ?? [], new Map(), getKey, getChildren)
  );

  useControlledState<T[]>(
    itemsProp ?? initialItemsProp ?? [],
    itemsProp,
    (prev, curr) => areItemsEqual(prev, curr, getKey, getChildren),
    () => setTreeState(buildTree(itemsProp, new Map(), getKey, getChildren))
  );

  const [selectedKeys, setSelectedKeys] = useState<Set<V>>(
    () => new Set<V>(selectedKeysProp ?? defaultSelectedKeysProp ?? [])
  );

  useControlledState<Iterable<V>>(
    () => new Set<V>(selectedKeysProp ?? defaultSelectedKeysProp ?? []),
    selectedKeysProp,
    (prev, curr) => prev === curr,
    () => setSelectedKeys(new Set<V>(selectedKeysProp))
  );

  const stateRef = useRef({ tree, selectedKeys });
  stateRef.current = { tree, selectedKeys };

  const setState = (newTree: TreeDataState<T, V>): TreeDataState<T, V> => {
    setTreeState(newTree);
    stateRef.current = { ...stateRef.current, tree: newTree };
    return newTree;
  };

  const handleSelectedKeysChange = useStableCallback((keys: Set<V>): void => {
    setSelectedKeys(keys);
    stateRef.current = { ...stateRef.current, selectedKeys: keys };
    onChange?.([...keys]);
  });

  const api: TreeState<T, V> = {
    nodes: tree.nodes,
    selectedKeys: [...selectedKeys],
    setSelectedKeys: keys => handleSelectedKeysChange(new Set(keys)),
    getItem: useStableCallback(function getItem(key: V) {
      return stateRef.current.tree.nodeMap.get(key);
    }),
    getPosition: useStableCallback(function getPosition(key: V) {
      const node = stateRef.current.tree.nodeMap.get(key);

      if (!node) {
        return undefined;
      }

      const parentKey = node.parentKey ?? null;
      const parentNode = parentKey ? stateRef.current.tree.nodeMap.get(parentKey) : null;
      const index = parentNode
        ? (parentNode.children?.findIndex(child => child.key === key) ?? -1)
        : stateRef.current.tree.nodes.findIndex(n => n.key === key);

      return { parentKey, index };
    }),
    append: useStableCallback(function append(parentKey: V | null, values: T[]): TreeDataState<T, V> {
      if (parentKey == null) {
        const currentTree = stateRef.current.tree;
        const { nodes: newNodes, nodeMap: newMap } = buildTree(
          values,
          new Map(currentTree.nodeMap),
          getKey,
          getChildren,
          null
        );

        return setState({
          nodes: [...currentTree.nodes, ...newNodes],
          nodeMap: newMap,
        });
      } else {
        const currentTree = stateRef.current.tree;
        const parentNode = currentTree.nodeMap.get(parentKey);

        if (!parentNode) {
          return currentTree;
        }

        const { nodes: newNodes } = buildTree(values, new Map(), getKey, getChildren, parentKey);

        const newTree = updateTree(
          currentTree.nodes,
          parentKey,
          parentNode => ({
            key: parentNode.key,
            parentKey: parentNode.parentKey,
            value: parentNode.value,
            children: [...(parentNode.children ?? []), ...newNodes],
          }),
          currentTree.nodeMap
        );

        return setState(newTree);
      }
    }),
    replace: useStableCallback(function replace(parentKey: V | null, values: T[]): TreeDataState<T, V> {
      if (parentKey == null) {
        const { nodes: newNodes, nodeMap: finalMap } = buildTree(values, new Map(), getKey, getChildren, null);

        return setState({
          nodes: newNodes,
          nodeMap: finalMap,
        });
      } else {
        const currentTree = stateRef.current.tree;
        const parentNode = currentTree.nodeMap.get(parentKey);

        if (!parentNode) {
          return currentTree;
        }

        const { nodes: newNodes } = buildTree(values, new Map(), getKey, getChildren, parentKey);

        const newTree = updateTree(
          currentTree.nodes,
          parentKey,
          parentNode => ({
            key: parentNode.key,
            parentKey: parentNode.parentKey,
            value: parentNode.value,
            children: newNodes,
          }),
          currentTree.nodeMap
        );

        return setState(newTree);
      }
    }),
    remove: useStableCallback(function remove(...keys: V[]): TreeDataState<T, V> | undefined {
      if (keys.length === 0) {
        return undefined;
      }

      const currentTree = stateRef.current.tree;
      let newItems = currentTree.nodes;
      let prevMap = currentTree.nodeMap;
      let newTree: TreeDataState<T, V> | undefined;

      for (const key of keys) {
        newTree = updateTree(newItems, key, () => null, prevMap);
        prevMap = newTree.nodeMap;
        newItems = newTree.nodes;
      }

      setState(newTree!);

      const selection = new Set(stateRef.current.selectedKeys);

      for (const key of stateRef.current.selectedKeys) {
        if (!newTree?.nodeMap.has(key)) {
          selection.delete(key);
        }
      }

      handleSelectedKeysChange(selection);
      return newTree;
    }),
    move: useStableCallback(function move(key: V, toParentKey: V | null, index: number): TreeDataState<T, V> {
      const currentTree = stateRef.current.tree;
      const { nodes: originalNodes, nodeMap: originalMap } = currentTree;
      const node = originalMap.get(key);

      if (!node) {
        return currentTree;
      }

      const { nodes: newNodes, nodeMap: newMap } = updateTree(originalNodes, key, () => null, originalMap);

      const movedNode = {
        ...node,
        parentKey: toParentKey,
      };

      let newTree: TreeDataState<T, V>;

      if (toParentKey == null) {
        addNode(movedNode, newMap);
        newTree = { nodes: [...newNodes.slice(0, index), movedNode, ...newNodes.slice(index)], nodeMap: newMap };
      } else {
        newTree = updateTree(
          newNodes,
          toParentKey,
          parentNode => ({
            key: parentNode.key,
            parentKey: parentNode.parentKey,
            value: parentNode.value,
            children: [
              ...(parentNode.children?.slice(0, index) ?? []),
              movedNode,
              ...(parentNode.children?.slice(index) ?? []),
            ],
          }),
          newMap
        );
      }

      return setState(newTree);
    }),
    moveBefore: useStableCallback(function moveBefore(key: V, keys: V[]): TreeDataState<T, V> {
      const currentTree = stateRef.current.tree;
      const { nodes, nodeMap } = currentTree;
      const node = nodeMap.get(key);

      if (!node) {
        return currentTree;
      }

      const toParentKey = node.parentKey ?? null;
      let parent: null | TreeNode<T, V> = null;
      if (toParentKey != null) {
        parent = nodeMap.get(toParentKey) ?? null;
      }
      const toIndex = parent?.children ? parent.children.indexOf(node) : nodes.indexOf(node);
      const newTree = moveItems(currentTree, keys, parent, toIndex);

      return setState(newTree);
    }),
    moveAfter: useStableCallback(function moveAfter(key: V, keys: V[]): TreeDataState<T, V> {
      const currentTree = stateRef.current.tree;
      const { nodes: items, nodeMap } = currentTree;
      const node = nodeMap.get(key);

      if (!node) {
        return currentTree;
      }

      const toParentKey = node.parentKey ?? null;
      let parent: null | TreeNode<T, V> = null;
      if (toParentKey != null) {
        parent = nodeMap.get(toParentKey) ?? null;
      }
      let toIndex = parent?.children ? parent.children.indexOf(node) : items.indexOf(node);
      toIndex++;
      const newTree = moveItems(currentTree, keys, parent, toIndex);

      return setState(newTree);
    }),
    externalSync: useStableCallback(function externalSync(): TreeDataState<T, V> {
      const newTree = buildTree(itemsProp ?? initialItemsProp ?? [], new Map(), getKey, getChildren);
      return setState(newTree);
    }),
  };

  return api;
}

// Helper functions

function areItemsEqual<T extends object, V extends TreeItemKey = string>(
  prev: T[] | undefined,
  next: T[] | undefined,
  getKey: (item: T) => V,
  getChildren: (item: T) => T[]
): boolean {
  if (prev === next) return true;
  if (!prev || !next) return false;
  if (prev.length !== next.length) return false;

  for (const [index, prevItem] of prev.entries()) {
    const nextItem = next[index]!;

    if (getKey(prevItem) !== getKey(nextItem)) {
      return false;
    }

    const prevChildren = getChildren(prevItem);
    const nextChildren = getChildren(nextItem);

    if (!areItemsEqual(prevChildren, nextChildren, getKey, getChildren)) {
      return false;
    }
  }

  return true;
}

function buildTree<T extends object, V extends TreeItemKey = string>(
  initialItems: T[] | null = [],
  map: Map<V, TreeNode<T, V>>,
  getKey: (item: T) => V,
  getChildren: (item: T) => T[],
  parentKey?: V | null
): TreeDataState<T, V> {
  if (initialItems == null) {
    initialItems = [];
  }

  const nodes: TreeNode<T, V>[] = [];

  for (const item of initialItems) {
    const node: TreeNode<T, V> = {
      key: getKey(item),
      parentKey: parentKey ?? null,
      value: item,
      children: null,
    };

    if (map.has(node.key)) {
      continue;
    }

    map.set(node.key, node);
    nodes.push(node);

    if (node.parentKey != null) {
      const parentNode = map.get(node.parentKey);
      if (parentNode) {
        parentNode.children = [...(parentNode.children ?? []), node];
      }
    }

    node.children = buildTree(getChildren(item), map, getKey, getChildren, node.key).nodes;
  }

  return {
    nodes,
    nodeMap: map,
  };
}

function updateTree<T extends object, V extends TreeItemKey = string>(
  nodes: TreeNode<T, V>[],
  key: V | null,
  update: (node: TreeNode<T, V>) => TreeNode<T, V> | null,
  originalMap: Map<V, TreeNode<T, V>>
): TreeDataState<T, V> {
  let node = key == null ? null : originalMap.get(key);
  if (node == null) {
    return { nodes, nodeMap: originalMap };
  }
  const map = new Map<V, TreeNode<T, V>>(originalMap);

  let newNode = update(node);

  if (newNode == null) {
    deleteNode(node, map);
  } else {
    addNode(newNode, map);
  }

  while (node && node.parentKey) {
    const nextParent: TreeNode<T, V> = map.get(node.parentKey)!;
    const copy: TreeNode<T, V> = {
      key: nextParent.key,
      parentKey: nextParent.parentKey,
      value: nextParent.value,
      children: null,
    };

    let children = nextParent.children;
    if (newNode == null && children) {
      children = children.filter(c => c !== node);
    }

    copy.children =
      children?.map(child => {
        if (child === node) {
          return newNode!;
        }

        return child;
      }) ?? null;

    map.set(copy.key, copy);

    newNode = copy;
    node = nextParent;
  }

  if (newNode == null) {
    nodes = nodes.filter(c => c !== node);
  }

  return {
    nodes: nodes.map(item => {
      if (item === node) {
        return newNode!;
      }

      return item;
    }),
    nodeMap: map,
  };
}

function addNode<T extends object, V extends TreeItemKey = string>(node: TreeNode<T, V>, map: Map<V, TreeNode<T, V>>) {
  map.set(node.key, node);
  if (node.children) {
    for (const child of node.children) {
      addNode(child, map);
    }
  }
}

function deleteNode<T extends object, V extends TreeItemKey = string>(
  node: TreeNode<T, V>,
  map: Map<V, TreeNode<T, V>>
) {
  map.delete(node.key);
  if (node.children) {
    for (const child of node.children) {
      deleteNode(child, map);
    }
  }
}

function moveItems<T extends object, V extends TreeItemKey = string>(
  state: TreeDataState<T, V>,
  keys: Iterable<V>,
  toParent: TreeNode<T, V> | null,
  toIndex: number
): TreeDataState<T, V> {
  const { nodes, nodeMap } = state;

  let parent = toParent;
  const removeKeys = new Set(keys);

  while (parent?.parentKey != null) {
    if (removeKeys.has(parent.key)) {
      throw new Error('Cannot move an item to be a child of itself.');
    }
    parent = nodeMap.get(parent.parentKey!) ?? null;
  }

  const originalToIndex = toIndex;

  const keyArray = Array.isArray(keys) ? keys : [...keys];

  const inOrderKeys: Map<V, number> = new Map();
  const removedItems: Array<TreeNode<T, V>> = [];
  let newNodes = nodes;
  let newMap = nodeMap;
  let i = 0;

  function traversal(
    node: TreeNode<T, V> | null,
    {
      inorder,
      postorder,
    }: { inorder?: (node: TreeNode<T, V> | null) => void; postorder?: (node: TreeNode<T, V> | null) => void }
  ) {
    inorder?.(node);
    if (node != null) {
      for (const child of node.children ?? []) {
        traversal(child, { inorder, postorder });
        postorder?.(child);
      }
    }
  }

  function inorder(child: TreeNode<T, V> | null) {
    if (child && keyArray.includes(child.key)) {
      inOrderKeys.set(child.key, i++);
    }
  }

  function postorder(child: TreeNode<T, V> | null) {
    if (child && keyArray.includes(child.key)) {
      removedItems.push({ ...newMap.get(child.key)!, parentKey: toParent?.key ?? null });
      const { nodes: nextNodes, nodeMap: nextMap } = updateTree(newNodes, child.key, () => null, newMap);
      newNodes = nextNodes;
      newMap = nextMap;
    }

    if (
      child &&
      (child.parentKey === toParent || child.parentKey === toParent?.key) &&
      keyArray.includes(child.key) &&
      (toParent?.children ? toParent.children.indexOf(child) : nodes.indexOf(child)) < originalToIndex
    ) {
      toIndex--;
    }
  }

  traversal({ children: nodes } as TreeNode<T, V>, { inorder, postorder });

  const inOrderItems = removedItems.toSorted((a, b) => (inOrderKeys.get(a.key)! > inOrderKeys.get(b.key)! ? 1 : -1));

  if (!toParent || toParent.key == null) {
    inOrderItems.forEach(movedNode => {
      addNode(movedNode, newMap);
    });
    return { nodes: [...newNodes.slice(0, toIndex), ...inOrderItems, ...newNodes.slice(toIndex)], nodeMap: newMap };
  }

  return updateTree(
    newNodes,
    toParent.key,
    parentNode => ({
      key: parentNode.key,
      parentKey: parentNode.parentKey,
      value: parentNode.value,
      children: [
        ...(parentNode.children?.slice(0, toIndex) ?? []),
        ...inOrderItems,
        ...(parentNode.children?.slice(toIndex) ?? []),
      ],
    }),
    newMap
  );
}

export const collectKeys = <T extends object, V extends TreeItemKey>(nodes: TreeNode<T, V>[]): V[] => {
  const keys: V[] = [];
  for (const child of nodes) {
    keys.push(child.key);
    if (child.children) {
      keys.push(...collectKeys(child.children));
    }
  }
  return keys;
};