import { original, produce } from 'immer'
import {
  FullTree,
  GetNodeKeyFunction,
  GetTreeItemChildren,
  NodeData,
  SearchData,
  TreeIndex,
  TreeItem,
  TreeNode,
  TreePath,
} from '../types'

export type WalkAndMapFunctionParameters = FullTree & {
  getNodeKey: GetNodeKeyFunction
  callback: Function
  ignoreCollapsed?: boolean
}

export interface FlatDataItem extends TreeNode, TreePath {
  lowerSiblingCounts: number[]
  parentNode: TreeItem
}

type NodeDataResult = {
  node?: TreeItem
  lowerSiblingCounts?: number[]
  path?: Array<string | number>
  nextIndex?: number
}

/**
 * Performs a depth-first traversal over all of the node descendants,
 * incrementing currentIndex by 1 for each
 */
const getNodeDataAtTreeIndexOrNextIndex = ({
  targetIndex,
  node,
  currentIndex,
  getNodeKey,
  path = [],
  lowerSiblingCounts = [],
  ignoreCollapsed = true,
  isPseudoRoot = false,
}: {
  targetIndex: number
  node: TreeItem
  currentIndex: number
  getNodeKey: GetNodeKeyFunction
  path?: Array<string | number>
  lowerSiblingCounts?: number[]
  ignoreCollapsed?: boolean
  isPseudoRoot?: boolean
}): NodeDataResult => {
  // The pseudo-root is not considered in the path
  const selfPath = isPseudoRoot
    ? []
    : [...path, getNodeKey({ node, treeIndex: currentIndex })]

  // Return target node when found
  if (currentIndex === targetIndex) {
    return {
      node,
      lowerSiblingCounts,
      path: selfPath,
    }
  }

  // Add one and continue for nodes with no children or hidden children
  if (
    !node?.children ||
    typeof node.children === 'function' ||
    (ignoreCollapsed && node?.expanded !== true)
  ) {
    return { nextIndex: currentIndex + 1 }
  }

  // Iterate over each child and their descendants and return the
  // target node if childIndex reaches the targetIndex
  let childIndex = currentIndex + 1
  const childCount = node.children.length
  for (let i = 0; i < childCount; i += 1) {
    const result = getNodeDataAtTreeIndexOrNextIndex({
      ignoreCollapsed,
      getNodeKey,
      targetIndex,
      node: node.children[i],
      currentIndex: childIndex,
      lowerSiblingCounts: [...lowerSiblingCounts, childCount - i - 1],
      path: selfPath,
    })

    if (result.node) {
      return result
    }

    childIndex = result.nextIndex!
  }

  // If the target node is not found, return the farthest traversed index
  return { nextIndex: childIndex }
}

export const getDescendantCount = ({
  node,
  ignoreCollapsed = true,
}: TreeNode & { ignoreCollapsed?: boolean }): number => {
  return (
    getNodeDataAtTreeIndexOrNextIndex({
      getNodeKey: () => 0,
      ignoreCollapsed,
      node,
      currentIndex: 0,
      targetIndex: -1,
    }).nextIndex! - 1
  )
}

const walkDescendants = ({
  callback,
  getNodeKey,
  ignoreCollapsed,
  isPseudoRoot = false,
  node,
  parentNode = undefined,
  currentIndex,
  path = [],
  lowerSiblingCounts = [],
}: {
  callback: Function
  getNodeKey: GetNodeKeyFunction
  ignoreCollapsed: boolean
  isPseudoRoot?: boolean
  node: TreeItem
  parentNode?: TreeItem
  currentIndex: number
  path?: Array<string | number>
  lowerSiblingCounts?: number[]
}): number | false => {
  // The pseudo-root is not considered in the path
  const selfPath = isPseudoRoot
    ? []
    : [...path, getNodeKey({ node, treeIndex: currentIndex })]
  const selfInfo = isPseudoRoot
    ? undefined
    : {
        node,
        parentNode,
        path: selfPath,
        lowerSiblingCounts,
        treeIndex: currentIndex,
      }

  if (!isPseudoRoot) {
    const callbackResult = callback(selfInfo)

    // Cut walk short if the callback returned false
    if (callbackResult === false) {
      return false
    }
  }

  // Return self on nodes with no children or hidden children
  if (
    !node.children ||
    (node.expanded !== true && ignoreCollapsed && !isPseudoRoot)
  ) {
    return currentIndex
  }

  // Get all descendants
  let childIndex = currentIndex
  const childCount = node.children.length
  if (typeof node.children !== 'function') {
    for (let i = 0; i < childCount; i += 1) {
      childIndex = walkDescendants({
        callback,
        getNodeKey,
        ignoreCollapsed,
        node: node.children[i],
        parentNode: isPseudoRoot ? undefined : node,
        currentIndex: childIndex + 1,
        lowerSiblingCounts: [...lowerSiblingCounts, childCount - i - 1],
        path: selfPath,
      }) as number

      // Cut walk short if the callback returned false
      if (childIndex === (false as any)) {
        return false
      }
    }
  }

  return childIndex
}

const mapDescendants = ({
  callback,
  getNodeKey,
  ignoreCollapsed,
  isPseudoRoot = false,
  node,
  parentNode = undefined,
  currentIndex,
  path = [],
  lowerSiblingCounts = [],
}: {
  callback: Function
  getNodeKey: GetNodeKeyFunction
  ignoreCollapsed: boolean
  isPseudoRoot?: boolean
  node: TreeItem
  parentNode?: TreeItem
  currentIndex: number
  path?: Array<string | number>
  lowerSiblingCounts?: number[]
}): { node: TreeItem; treeIndex: number } => {
  const nextNode = { ...node }

  // The pseudo-root is not considered in the path
  const selfPath = isPseudoRoot
    ? []
    : [...path, getNodeKey({ node: nextNode, treeIndex: currentIndex })]
  const selfInfo = {
    node: nextNode,
    parentNode,
    path: selfPath,
    lowerSiblingCounts,
    treeIndex: currentIndex,
  }

  // Return self on nodes with no children or hidden children
  if (
    !nextNode.children ||
    (nextNode.expanded !== true && ignoreCollapsed && !isPseudoRoot)
  ) {
    return {
      treeIndex: currentIndex,
      node: callback(selfInfo),
    }
  }

  // Get all descendants
  let childIndex = currentIndex
  const childCount = nextNode.children.length
  if (typeof nextNode.children !== 'function') {
    nextNode.children = nextNode.children.map((child: TreeItem, i: number) => {
      const mapResult = mapDescendants({
        callback,
        getNodeKey,
        ignoreCollapsed,
        node: child,
        parentNode: isPseudoRoot ? undefined : nextNode,
        currentIndex: childIndex + 1,
        lowerSiblingCounts: [...lowerSiblingCounts, childCount - i - 1],
        path: selfPath,
      })
      childIndex = mapResult.treeIndex

      return mapResult.node
    })
  }

  return {
    node: callback(selfInfo),
    treeIndex: childIndex,
  }
}

export const getVisibleNodeCount = ({ treeData }: FullTree): number => {
  const traverse = (node: TreeItem): number => {
    if (
      !node.children ||
      node.expanded !== true ||
      typeof node.children === 'function'
    ) {
      return 1
    }

    return (
      1 +
      node.children.reduce(
        (total: number, currentNode: TreeItem) => total + traverse(currentNode),
        0
      )
    )
  }

  return treeData.reduce(
    (total, currentNode) => total + traverse(currentNode),
    0
  )
}

export const getVisibleNodeInfoAtIndex = ({
  treeData,
  index: targetIndex,
  getNodeKey,
}: FullTree & {
  index: number
  getNodeKey: GetNodeKeyFunction
}): (TreeNode & TreePath & { lowerSiblingCounts: number[] }) | null => {
  if (!treeData || treeData.length === 0) {
    return null
  }

  // Call the tree traversal with a pseudo-root node
  const result = getNodeDataAtTreeIndexOrNextIndex({
    targetIndex,
    getNodeKey,
    node: {
      children: treeData,
      expanded: true,
    },
    currentIndex: -1,
    path: [],
    lowerSiblingCounts: [],
    ignoreCollapsed: true,
    isPseudoRoot: true,
  })

  if (result.node) {
    return result as TreeNode & TreePath & { lowerSiblingCounts: number[] }
  }

  return null
}

export const walk = ({
  treeData,
  getNodeKey,
  callback,
  ignoreCollapsed = true,
}: WalkAndMapFunctionParameters): void => {
  if (!treeData || treeData.length === 0) {
    return
  }

  walkDescendants({
    callback,
    getNodeKey,
    ignoreCollapsed,
    isPseudoRoot: true,
    node: { children: treeData },
    currentIndex: -1,
    path: [],
    lowerSiblingCounts: [],
  })
}

export const map = ({
  treeData,
  getNodeKey,
  callback,
  ignoreCollapsed = true,
}: WalkAndMapFunctionParameters): TreeItem[] => {
  if (!treeData || treeData.length === 0) {
    return []
  }

  return mapDescendants({
    callback,
    getNodeKey,
    ignoreCollapsed,
    isPseudoRoot: true,
    node: { children: treeData },
    currentIndex: -1,
    path: [],
    lowerSiblingCounts: [],
  }).node.children as TreeItem[]
}

export const toggleExpandedForAll = ({
  treeData,
  expanded = true,
}: FullTree & {
  expanded?: boolean
}): TreeItem[] => {
  return map({
    treeData,
    callback: ({ node }: { node: TreeItem }) => ({ ...node, expanded }),
    getNodeKey: ({ treeIndex }: TreeIndex) => treeIndex,
    ignoreCollapsed: false,
  })
}

export const changeNodeAtPath = ({
  treeData,
  path,
  newNode,
  getNodeKey,
  ignoreCollapsed = true,
}: FullTree &
  TreePath & {
    newNode: Function | any
    getNodeKey: GetNodeKeyFunction
    ignoreCollapsed?: boolean
  }): TreeItem[] => {
  if (!treeData || treeData.length === 0) return []

  return produce(treeData, (draft) => {
    let currentNode = { children: draft } as any // Pseudo-root
    let currentTreeIndex = -1

    for (const [i, key] of path.entries()) {
      const isLast = i === path.length - 1

      if (!currentNode.children) {
        throw new Error('Path referenced children of node with no children.')
      }

      // Find the child matching the key
      let foundIndex = -1
      for (let j = 0; j < currentNode.children.length; j++) {
        const child = currentNode.children[j]
        const childIndex = currentTreeIndex + 1

        // We need to calculate indices to match keys, but we don't need
        // to modify the index logic significantly for Immer, just navigation
        if (getNodeKey({ node: child, treeIndex: childIndex }) === key) {
          foundIndex = j
          currentTreeIndex = childIndex
          break
        }

        // Increment index by skipping descendants
        currentTreeIndex +=
          1 + getDescendantCount({ node: child, ignoreCollapsed })
      }

      if (foundIndex === -1) {
        throw new Error('No node found at the given path.')
      }

      if (isLast) {
        const targetNode = currentNode.children[foundIndex]
        const result =
          typeof newNode === 'function'
            ? newNode({ node: targetNode, treeIndex: currentTreeIndex })
            : newNode

        if (result === undefined || result === null) {
          currentNode.children.splice(foundIndex, 1)
        } else {
          currentNode.children[foundIndex] = result
        }
      } else {
        currentNode = currentNode.children[foundIndex]
      }
    }
  })
}

export const removeNodeAtPath = ({
  treeData,
  path,
  getNodeKey,
  ignoreCollapsed = true,
}: FullTree &
  TreePath & {
    getNodeKey: GetNodeKeyFunction
    ignoreCollapsed?: boolean
  }): TreeItem[] => {
  return changeNodeAtPath({
    treeData,
    path,
    getNodeKey,
    ignoreCollapsed,
    newNode: undefined, // Delete the node
  })
}

export const removeNode = ({
  treeData,
  path,
  getNodeKey,
  ignoreCollapsed = true,
}: FullTree &
  TreePath & {
    getNodeKey: GetNodeKeyFunction
    ignoreCollapsed?: boolean
  }): (FullTree & TreeNode & TreeIndex) | undefined => {
  let removedNode: TreeItem | undefined
  let removedTreeIndex: number | undefined
  const nextTreeData = changeNodeAtPath({
    treeData,
    path,
    getNodeKey,
    ignoreCollapsed,
    newNode: ({ node, treeIndex }: { node: TreeItem; treeIndex: number }) => {
      removedNode = original(node) || node

      removedTreeIndex = treeIndex

      return undefined
    },
  })

  return {
    treeData: nextTreeData,
    node: removedNode!,
    treeIndex: removedTreeIndex!,
  }
}

export const getNodeAtPath = ({
  treeData,
  path,
  getNodeKey,
  ignoreCollapsed = true,
}: FullTree &
  TreePath & {
    getNodeKey: GetNodeKeyFunction
    ignoreCollapsed?: boolean
  }): (TreeNode & TreeIndex) | null => {
  let foundNodeInfo: (TreeNode & TreeIndex) | undefined

  try {
    changeNodeAtPath({
      treeData,
      path,
      getNodeKey,
      ignoreCollapsed,
      newNode: ({ node, treeIndex }: GetTreeItemChildren) => {
        foundNodeInfo = { node: original(node) || node, treeIndex }
        return node
      },
    })
  } catch {
    // Ignore the error -- the null return will be explanation enough
  }

  return foundNodeInfo ?? null
}

export const addNodeUnderParent = ({
  treeData,
  newNode,
  parentKey = undefined,
  getNodeKey,
  ignoreCollapsed = true,
  expandParent = false,
  addAsFirstChild = false,
}: FullTree & {
  newNode: TreeItem
  parentKey: number | string | undefined | null
  getNodeKey: GetNodeKeyFunction
  ignoreCollapsed?: boolean
  expandParent?: boolean
  addAsFirstChild?: boolean
}): FullTree & TreeIndex => {
  if (parentKey === null || parentKey === undefined) {
    const newTreeData = addAsFirstChild
      ? [newNode, ...(treeData || [])]
      : [...(treeData || []), newNode]
    return {
      treeData: newTreeData,
      treeIndex: addAsFirstChild ? 0 : (treeData || []).length,
    }
  }

  let insertedTreeIndex = -1
  let found = false

  const nextTreeData = produce(treeData, (draft) => {
    // Helper to find parent and insert
    const findAndInsert = (
      nodes: TreeItem[],
      currentTreeIndex: number
    ): number => {
      let indexCounter = currentTreeIndex

      for (const node of nodes) {
        // Calculate key
        const key = getNodeKey({ node, treeIndex: indexCounter })

        if (key === parentKey) {
          found = true
          if (expandParent) {
            node.expanded = true
          }

          if (!node.children) {
            node.children = [newNode]
            insertedTreeIndex = indexCounter + 1
            return -1 // Stop traversal, we found it
          }

          if (typeof node.children === 'function') {
            throw new TypeError('Cannot add to children defined by a function')
          }

          // Calculate where the new node lands in the index
          let childIndexOffset = indexCounter + 1
          if (!addAsFirstChild) {
            for (const child of node.children) {
              childIndexOffset +=
                1 + getDescendantCount({ node: child, ignoreCollapsed })
            }
          }
          insertedTreeIndex = childIndexOffset

          if (addAsFirstChild) {
            node.children.unshift(newNode)
          } else {
            node.children.push(newNode)
          }
          return -1 // Stop traversal
        }

        // Standard traversal increment
        // If not found, add self (1) + descendants
        const descendants = getDescendantCount({ node, ignoreCollapsed })
        const nextIndex = indexCounter + 1 + descendants

        // If the node has children, we might need to look inside them
        if (
          node.children &&
          typeof node.children !== 'function' &&
          (node.expanded || !ignoreCollapsed)
        ) {
          // If we are strictly counting indices, we can skip traversing children
          // if we know the key isn't in there?
          // No, getNodeKey depends on treeIndex, so we must traverse to calculate index accurately.
          const result = findAndInsert(node.children, indexCounter + 1)
          if (result === -1) return -1
        }

        indexCounter = nextIndex
      }
      return indexCounter
    }

    findAndInsert(draft, 0)
  })

  if (!found) {
    throw new Error('No node found with the given key.')
  }

  return {
    treeData: nextTreeData,
    treeIndex: insertedTreeIndex,
  }
}

interface AddNodeResult {
  node: TreeItem
  nextIndex: number
  insertedTreeIndex?: number
  parentPath?: Array<string | number>
  parentNode?: TreeItem
}

const addNodeAtDepthAndIndex = ({
  targetDepth,
  minimumTreeIndex,
  newNode,
  ignoreCollapsed,
  expandParent,
  isPseudoRoot = false,
  isLastChild,
  node,
  currentIndex,
  currentDepth,
  getNodeKey,
  path = [],
}: {
  targetDepth: number
  minimumTreeIndex: number
  newNode: TreeItem
  ignoreCollapsed: boolean
  expandParent: boolean
  isPseudoRoot?: boolean
  isLastChild: boolean
  node: TreeItem
  currentIndex: number
  currentDepth: number
  getNodeKey: GetNodeKeyFunction
  path?: Array<string | number>
}): AddNodeResult => {
  const selfPath = (n: TreeItem) =>
    isPseudoRoot
      ? []
      : [...path, getNodeKey({ node: n, treeIndex: currentIndex })]

  // If the current position is the only possible place to add, add it
  if (
    currentIndex >= minimumTreeIndex - 1 ||
    (isLastChild && !(node.children && node.children.length > 0))
  ) {
    if (typeof node.children === 'function') {
      throw new TypeError('Cannot add to children defined by a function')
    } else {
      const extraNodeProps = expandParent ? { expanded: true } : {}
      const nextNode = {
        ...node,

        ...extraNodeProps,
        children: node.children ? [newNode, ...node.children] : [newNode],
      }

      return {
        node: nextNode,
        nextIndex: currentIndex + 2,
        insertedTreeIndex: currentIndex + 1,
        parentPath: selfPath(nextNode),
        parentNode: isPseudoRoot ? undefined : nextNode,
      }
    }
  }

  // If this is the target depth for the insertion,
  // i.e., where the newNode can be added to the current node's children
  if (currentDepth >= targetDepth - 1) {
    // Skip over nodes with no children or hidden children
    if (
      !node.children ||
      typeof node.children === 'function' ||
      (node.expanded !== true && ignoreCollapsed && !isPseudoRoot)
    ) {
      return { node, nextIndex: currentIndex + 1 }
    }

    // Scan over the children to see if there's a place among them that fulfills
    // the minimumTreeIndex requirement
    let childIndex = currentIndex + 1
    let insertedTreeIndex: number | undefined
    let insertIndex: number | undefined
    for (let i = 0; i < node.children.length; i += 1) {
      // If a valid location is found, mark it as the insertion location and
      // break out of the loop
      if (childIndex >= minimumTreeIndex) {
        insertedTreeIndex = childIndex
        insertIndex = i
        break
      }

      // Increment the index by the child itself plus the number of descendants it has
      childIndex +=
        1 + getDescendantCount({ node: node.children[i], ignoreCollapsed })
    }

    // If no valid indices to add the node were found
    if (insertIndex === null || insertIndex === undefined) {
      // If the last position in this node's children is less than the minimum index
      // and there are more children on the level of this node, return without insertion
      if (childIndex < minimumTreeIndex && !isLastChild) {
        return { node, nextIndex: childIndex }
      }

      // Use the last position in the children array to insert the newNode
      insertedTreeIndex = childIndex
      insertIndex = node.children.length
    }

    // Insert the newNode at the insertIndex
    const nextNode = {
      ...node,
      children: node.children.toSpliced(insertIndex, 0, newNode),
    }

    // Return node with successful insert result
    return {
      node: nextNode,
      nextIndex: childIndex,
      insertedTreeIndex,
      parentPath: selfPath(nextNode),
      parentNode: isPseudoRoot ? undefined : nextNode,
    }
  }

  // Skip over nodes with no children or hidden children
  if (
    !node.children ||
    typeof node.children === 'function' ||
    (node.expanded !== true && ignoreCollapsed && !isPseudoRoot)
  ) {
    return { node, nextIndex: currentIndex + 1 }
  }

  // Get all descendants
  let insertedTreeIndex: number | undefined
  let pathFragment: Array<string | number> | undefined
  let parentNode: TreeItem | undefined
  let childIndex = currentIndex + 1
  let newChildren = node.children
  if (typeof newChildren !== 'function') {
    newChildren = newChildren.map((child: TreeItem, i: number) => {
      if (insertedTreeIndex !== null && insertedTreeIndex !== undefined) {
        return child
      }

      const mapResult = addNodeAtDepthAndIndex({
        targetDepth,
        minimumTreeIndex,
        newNode,
        ignoreCollapsed,
        expandParent,
        isLastChild: isLastChild && i === newChildren.length - 1,
        node: child,
        currentIndex: childIndex,
        currentDepth: currentDepth + 1,
        getNodeKey,
        path: [], // Cannot determine the parent path until the children have been processed
      })

      if (
        'insertedTreeIndex' in mapResult &&
        mapResult.insertedTreeIndex !== undefined
      ) {
        ;({
          insertedTreeIndex,
          parentNode,
          parentPath: pathFragment,
        } = mapResult)
      }

      childIndex = mapResult.nextIndex

      return mapResult.node
    })
  }

  const nextNode = { ...node, children: newChildren }
  const result: AddNodeResult = {
    node: nextNode,
    nextIndex: childIndex,
  }

  if (insertedTreeIndex !== null && insertedTreeIndex !== undefined) {
    result.insertedTreeIndex = insertedTreeIndex
    result.parentPath = [...selfPath(nextNode), ...(pathFragment || [])]
    result.parentNode = parentNode
  }

  return result
}

export const insertNode = ({
  treeData,
  depth: targetDepth,
  minimumTreeIndex,
  newNode,
  getNodeKey,
  ignoreCollapsed = true,
  expandParent = false,
}: FullTree & {
  depth: number
  newNode: TreeItem
  minimumTreeIndex: number
  ignoreCollapsed?: boolean
  expandParent?: boolean
  getNodeKey: GetNodeKeyFunction
}): FullTree & TreeIndex & TreePath & { parentNode: TreeItem | null } => {
  if (!treeData && targetDepth === 0) {
    return {
      treeData: [newNode],
      treeIndex: 0,
      path: [getNodeKey({ node: newNode, treeIndex: 0 }) as number],
      parentNode: null,
    }
  }

  const insertResult = addNodeAtDepthAndIndex({
    targetDepth,
    minimumTreeIndex,
    newNode,
    ignoreCollapsed,
    expandParent,
    getNodeKey,
    isPseudoRoot: true,
    isLastChild: true,
    node: { children: treeData },
    currentIndex: -1,
    currentDepth: -1,
  })

  if (
    !('insertedTreeIndex' in insertResult) ||
    insertResult.insertedTreeIndex === undefined
  ) {
    throw new Error('No suitable position found to insert.')
  }

  const treeIndex = insertResult.insertedTreeIndex
  return {
    treeData: insertResult.node.children as TreeItem[],
    treeIndex,
    path: [
      ...(insertResult.parentPath || []),
      getNodeKey({ node: newNode, treeIndex }),
    ] as number[],
    parentNode: insertResult.parentNode ?? null,
  }
}

export const getFlatDataFromTree = ({
  treeData,
  getNodeKey,
  ignoreCollapsed = true,
}: FullTree & {
  getNodeKey: GetNodeKeyFunction
  ignoreCollapsed?: boolean
}): FlatDataItem[] => {
  if (!treeData || treeData.length === 0) {
    return []
  }

  const flattened: FlatDataItem[] = []
  walk({
    treeData,
    getNodeKey,
    ignoreCollapsed,
    callback: (nodeInfo: FlatDataItem) => {
      flattened.push(nodeInfo)
    },
  })

  return flattened
}

export const getTreeFromFlatData = ({
  flatData,
  getKey = (node) => node.id,
  getParentKey = (node) => node.parentId,
  rootKey = '0',
}: {
  flatData: any
  getKey: (node: any) => string
  getParentKey: (node: any) => string
  rootKey: string | null
}) => {
  if (!flatData) {
    return []
  }

  const childrenToParents = Object.groupBy(flatData, (child: any) =>
    getParentKey(child)
  )

  if (rootKey === null || !childrenToParents[rootKey]) {
    return []
  }

  const trav = (parent: any): any => {
    const parentKey = getKey(parent)
    const children = childrenToParents[parentKey]
    if (children) {
      return {
        ...parent,
        children: children.map((child: any) => trav(child)),
      }
    }

    return { ...parent }
  }

  return childrenToParents[rootKey]!.map((child: any) => trav(child))
}

export const isDescendant = (older: TreeItem, younger: TreeItem): boolean => {
  return (
    !!older.children &&
    typeof older.children !== 'function' &&
    older.children.some(
      (child) => child === younger || isDescendant(child, younger)
    )
  )
}

export const getDepth = (node: TreeItem, depth = 0): number => {
  if (!node.children) {
    return depth
  }

  if (typeof node.children === 'function') {
    return depth + 1
  }

  return node.children.reduce(
    (deepest, child) => Math.max(deepest, getDepth(child, depth + 1)),
    depth
  )
}

export const find = ({
  getNodeKey,
  treeData,
  searchQuery,
  searchMethod,
  searchFocusOffset,
  expandAllMatchPaths = false,
  expandFocusMatchPaths = true,
}: FullTree & {
  getNodeKey: GetNodeKeyFunction
  searchQuery?: string | number
  searchMethod: (data: SearchData) => boolean
  searchFocusOffset?: number
  expandAllMatchPaths?: boolean
  expandFocusMatchPaths?: boolean
}): { matches: NodeData[] } & FullTree => {
  let matchCount = 0
  const trav = ({
    isPseudoRoot = false,
    node,
    currentIndex,
    path = [],
  }: {
    isPseudoRoot?: boolean
    node: TreeItem
    currentIndex: number
    path?: Array<string | number>
  }) => {
    let matches: any[] = []
    let isSelfMatch = false
    let hasFocusMatch = false
    // The pseudo-root is not considered in the path
    const selfPath = isPseudoRoot
      ? []
      : [...path, getNodeKey({ node, treeIndex: currentIndex })]
    const extraInfo = isPseudoRoot
      ? undefined
      : {
          path: selfPath,
          treeIndex: currentIndex,
        }

    // Nodes with with children that aren't lazy
    const hasChildren =
      node.children &&
      typeof node.children !== 'function' &&
      node.children.length > 0

    // Examine the current node to see if it is a match
    if (
      !isPseudoRoot &&
      searchMethod({
        ...extraInfo,
        node,
        searchQuery: searchQuery as string,
      } as SearchData)
    ) {
      if (matchCount === searchFocusOffset) {
        hasFocusMatch = true
      }

      // Keep track of the number of matching nodes, so we know when the searchFocusOffset
      //  is reached
      matchCount += 1

      // We cannot add this node to the matches right away, as it may be changed
      //  during the search of the descendants. The entire node is used in
      //  comparisons between nodes inside the `matches` and `treeData` results
      //  of this method (`find`)
      isSelfMatch = true
    }

    let childIndex = currentIndex
    const newNode = { ...node }
    if (hasChildren) {
      // Get all descendants
      newNode.children = (newNode.children as TreeItem[]).map(
        (child: TreeItem) => {
          const mapResult = trav({
            node: child,
            currentIndex: childIndex + 1,
            path: selfPath,
          })

          // Ignore hidden nodes by only advancing the index counter to the returned treeIndex
          // if the child is expanded.
          //
          // The child could have been expanded from the start,
          // or expanded due to a matching node being found in its descendants
          if (mapResult.node.expanded) {
            childIndex = mapResult.treeIndex
          } else {
            childIndex += 1
          }

          if (mapResult.matches.length > 0 || mapResult.hasFocusMatch) {
            matches = [...matches, ...mapResult.matches]
            if (mapResult.hasFocusMatch) {
              hasFocusMatch = true
            }

            // Expand the current node if it has descendants matching the search
            // and the settings are set to do so.
            if (
              (expandAllMatchPaths && mapResult.matches.length > 0) ||
              ((expandAllMatchPaths || expandFocusMatchPaths) &&
                mapResult.hasFocusMatch)
            ) {
              newNode.expanded = true
            }
          }

          return mapResult.node
        }
      )
    }

    // Cannot assign a treeIndex to hidden nodes
    if (!isPseudoRoot && !newNode.expanded) {
      matches = matches.map((match) => ({
        ...match,
        treeIndex: undefined,
      }))
    }

    // Add this node to the matches if it fits the search criteria.
    // This is performed at the last minute so newNode can be sent in its final form.
    if (isSelfMatch) {
      matches = [{ ...extraInfo, node: newNode }, ...matches]
    }

    return {
      node: matches.length > 0 ? newNode : node,
      matches,
      hasFocusMatch,
      treeIndex: childIndex,
    }
  }

  const result = trav({
    node: { children: treeData },
    isPseudoRoot: true,
    currentIndex: -1,
  })

  return {
    matches: result.matches,
    treeData: result.node.children as TreeItem[],
  }
}