export type HasParent = { id: number; parent?: { id: number } | null };
export type TreeNode<T extends HasParent> = T & { children: Array<TreeNode<T>>; expanded: boolean };
export type RootNode<T extends HasParent> = { id?: number; children: Array<TreeNode<T>> };

export function arrayToTree<T extends HasParent>(nodes: T[], currentState?: RootNode<T>): RootNode<T> {
    const topLevelNodes: Array<TreeNode<T>> = [];
    const mappedArr: { [id: string]: TreeNode<T> } = {};
    const currentStateMap = treeToMap(currentState);

    // First map the nodes of the array to an object -> create a hash table.
    for (const node of nodes) {
        mappedArr[node.id] = { ...(node as any), children: [] };
    }

    for (const id of nodes.map(n => n.id)) {
        if (mappedArr.hasOwnProperty(id)) {
            const mappedElem = mappedArr[id];
            mappedElem.expanded = currentStateMap.get(id)?.expanded ?? false;
            const parent = mappedElem.parent;
            if (!parent) {
                continue;
            }
            // If the element is not at the root level, add it to its parent array of children.
            const parentIsRoot = !mappedArr[parent.id];
            if (!parentIsRoot) {
                if (mappedArr[parent.id]) {
                    mappedArr[parent.id].children.push(mappedElem);
                } else {
                    mappedArr[parent.id] = { children: [mappedElem] } as any;
                }
            } else {
                topLevelNodes.push(mappedElem);
            }
        }
    }
    // tslint:disable-next-line:no-non-null-assertion
    const rootId = topLevelNodes.length ? topLevelNodes[0].parent!.id : undefined;
    return { id: rootId, children: topLevelNodes };
}

/**
 * Converts an existing tree (as generated by the arrayToTree function) into a flat
 * Map. This is used to persist certain states (e.g. `expanded`) when re-building the
 * tree.
 */
function treeToMap<T extends HasParent>(tree?: RootNode<T>): Map<number, TreeNode<T>> {
    const nodeMap = new Map<number, TreeNode<T>>();
    function visit(node: TreeNode<T>) {
        nodeMap.set(node.id, node);
        node.children.forEach(visit);
    }
    if (tree) {
        visit(tree as TreeNode<T>);
    }
    return nodeMap;
}