import type { MarkdownHeading } from "astro";

export type HeadingsTree = Array<MarkdownHeading | HeadingsTree>;

// could take a list like [5, 1, 2, 2, 3, 5, 1] and should return
// [ [ [ [ [ 5 ] ] ] ], 1, [ 2, 2, [3, [[5]]]], 1]
export function makeHeadingsTree(headings: MarkdownHeading[]): HeadingsTree {
  let index = 0;
  let currDepth = 1;
  const treeStack = [[]] as Array<HeadingsTree>;

  while (index < headings.length) {
    const currHeading = headings[index];

    // If the next heading is at the current depth, push it into the tree
    if (currHeading.depth == currDepth) {
      treeStack[treeStack.length - 1].push(currHeading);
      index += 1;
    }

    // If the next heading is deeper, go deeper
    if (currHeading.depth > currDepth) {
      const nextLevel = [] as HeadingsTree;
      treeStack[treeStack.length - 1].push(nextLevel);
      treeStack.push(nextLevel);
      currDepth += 1;
    }

    // If the next heading is shallower, back out
    if (currHeading.depth < currDepth) {
      treeStack.pop();
      currDepth -= 1;
    }
  }

  return treeStack[0];
}

export function removeEmptyLevels(tree: HeadingsTree): HeadingsTree {
  const newLevel = tree.map((node) =>
    Array.isArray(node) ? removeEmptyLevels(node) : node,
  );

  if (tree.every(Array.isArray)) {
    return tree.flatMap((node) => node);
  }

  return tree;
}
