import type { Key, ReactNode } from 'react';

/**
 * The tree's data shape and the pure walks over it. No React, no state — which
 * is why this file exists: `use-tree-keyboard.ts` needs `isBranch` and the
 * flattened row shape, and importing them from `tree.tsx` while `tree.tsx`
 * imports the hook would put the two modules in a cycle. Both import this
 * instead, and nothing here imports either of them.
 */

export interface TreeNode {
  key: Key;
  title: ReactNode;
  children?: TreeNode[];
  disabled?: boolean;
  /** Exclude just the checkbox, leaving the node selectable. */
  disableCheckbox?: boolean;
  /** Force leaf rendering for a branch whose children load lazily. */
  isLeaf?: boolean;
}

/** A node plus where it sits, which is what rendering and navigation both need. */
export interface FlatNode {
  node: TreeNode;
  depth: number;
  parentKey: Key | null;
}

export const isBranch = (node: TreeNode) => !node.isLeaf && !!node.children?.length;

/** Walks the tree, emitting only what is currently visible. */
export const flatten = (
  nodes: TreeNode[],
  expanded: Set<Key>,
  depth = 0,
  parentKey: Key | null = null,
  out: FlatNode[] = []
) => {
  for (const node of nodes) {
    out.push({ node, depth, parentKey });
    if (isBranch(node) && expanded.has(node.key)) {
      flatten(node.children!, expanded, depth + 1, node.key, out);
    }
  }
  return out;
};

export const branchKeys = (nodes: TreeNode[], out: Key[] = []) => {
  for (const node of nodes) {
    if (isBranch(node)) {
      out.push(node.key);
      branchKeys(node.children!, out);
    }
  }
  return out;
};

/**
 * Index of every node by key, plus its parent and the leaves beneath it.
 *
 * `leavesOf` is the reason this walk exists rather than a helper called per
 * node. Checkbox state is *derived* from a branch's leaves (see `checkStateOf`),
 * so a per-node `leafKeysUnder` re-walked the subtree once for every visible row
 * on every render, and once more for every branch on every toggle — in the one
 * component that offers `virtual` for trees of thousands of nodes. Collecting
 * the leaves on the way back up costs the same single walk the index already
 * does, and every later question is a `Map` lookup.
 *
 * A node with no children is its own only leaf, which is what the old helper
 * did too — `isLeaf` marks a branch whose children load later and deliberately
 * does not enter into it.
 */
export const indexTree = (nodes: TreeNode[]) => {
  const byKey = new Map<Key, TreeNode>();
  const parentOf = new Map<Key, Key>();
  const leavesOf = new Map<Key, Key[]>();

  const walk = (list: TreeNode[], parent: Key | null): Key[] => {
    const leaves: Key[] = [];

    for (const node of list) {
      byKey.set(node.key, node);
      if (parent !== null) parentOf.set(node.key, parent);

      const own = node.children?.length ? walk(node.children, node.key) : [node.key];
      leavesOf.set(node.key, own);
      leaves.push(...own);
    }

    return leaves;
  };

  walk(nodes, null);
  return { byKey, parentOf, leavesOf };
};