feuerwerk/src/core/model/hierarchy.ts

import { type Removed, removeWhere } from '../util/list';
import type { ObjectGroup, SceneObject } from './SceneObject';
import { descendantsBelow, parentChain } from './walk';

/**
 * `objects` stays a flat list in creation order; the tree is implied by `SceneObject.parentId`.
 * Anything holding such a list fits – a scene, or the template inside a container asset.
 */
export interface ObjectTree {
  objects: SceneObject[];
  /** Hierarchy folders for top-level objects (#241) – they order the tree but never shape it. */
  groups?: ObjectGroup[];
}

export function findObject(tree: ObjectTree, id: string): SceneObject | undefined {
  return tree.objects.find((o) => o.id === id);
}

function childrenOf(tree: ObjectTree, parentId: string): SceneObject[] {
  return tree.objects.filter((o) => o.parentId === parentId);
}

/** Children of `parentId` (`''` = top level) in sibling order – what a reorder rearranges. */
export function childIdsOf(tree: ObjectTree, parentId: string): string[] {
  return childrenOf(tree, parentId).map((o) => o.id);
}

function parentOf(tree: ObjectTree, object: SceneObject): SceneObject | undefined {
  return object.parentId ? findObject(tree, object.parentId) : undefined;
}

/** Depth-first, parents before children. */
export function descendantsOf(tree: ObjectTree, id: string): SceneObject[] {
  return descendantsBelow(
    id,
    (parentId) => childrenOf(tree, parentId),
    (o) => o.id,
  );
}

export function subtreeIds(tree: ObjectTree, rootIds: Iterable<string>): Set<string> {
  const ids = new Set<string>();
  for (const rootId of rootIds) {
    ids.add(rootId);
    for (const child of descendantsOf(tree, rootId)) ids.add(child.id);
  }
  return ids;
}

/** Returns the removed objects with their former indices (see `restoreAt`). */
export function removeSubtrees(tree: ObjectTree, rootIds: readonly string[]): Removed<SceneObject>[] {
  const ids = subtreeIds(tree, rootIds);
  return removeWhere(tree.objects, (o) => ids.has(o.id));
}

/** An object and every parent above it, nearest first. */
export function chainOf(tree: ObjectTree, object: SceneObject): SceneObject[] {
  return [...parentChain(object, (o) => parentOf(tree, o))];
}

/** Nearest parent first. */
export function ancestorIdsOf(tree: ObjectTree, id: string): string[] {
  const object = findObject(tree, id);
  return object
    ? chainOf(tree, object)
        .slice(1)
        .map((o) => o.id)
    : [];
}

export function isSelfOrAncestor(tree: ObjectTree, ancestorId: string, objectId: string): boolean {
  const object = findObject(tree, objectId);
  return object !== undefined && chainOf(tree, object).some((o) => o.id === ancestorId);
}

/** Of the given ids, those with no ancestor among the ids themselves. */
export function subtreeRoots(tree: ObjectTree, ids: readonly string[]): string[] {
  const set = new Set(ids);
  return ids.filter((id) => !ancestorIdsOf(tree, id).some((ancestorId) => set.has(ancestorId)));
}

export interface TreeEntry {
  object: SceneObject;
  depth: number;
}

/**
 * The stored list with the members of each group moved up behind the first of them: a group is
 * listed – and drawn at equal `z` – as one block wherever its first member stands (#241).
 */
export function groupedOrder(objects: readonly SceneObject[]): readonly SceneObject[] {
  const blocks = new Map<string, SceneObject[]>();
  for (const object of objects) {
    if (object.parentId || !object.groupId) continue;
    const block = blocks.get(object.groupId);
    if (block) block.push(object);
    else blocks.set(object.groupId, [object]);
  }
  if (blocks.size === 0) return objects;
  const result: SceneObject[] = [];
  for (const object of objects) {
    const block = object.parentId || !object.groupId ? undefined : blocks.get(object.groupId);
    if (!block) result.push(object);
    else if (block[0] === object) result.push(...block);
  }
  return result;
}

/**
 * Top-level objects in storage order, groups kept together, each followed by its subtree – the
 * order the hierarchy panel shows and the tie-break order for equal `z`.
 */
export function treeOrder(tree: ObjectTree): TreeEntry[] {
  const result: TreeEntry[] = [];
  const ordered = groupedOrder(tree.objects);
  const ids = new Set(tree.objects.map((o) => o.id));
  const visit = (parentId: string, depth: number) => {
    for (const object of ordered) {
      if (object.parentId !== parentId) continue;
      result.push({ object, depth });
      visit(object.id, depth + 1);
    }
  };
  visit('', 0);
  // Objects whose parent is missing are listed at the top level rather than dropped.
  for (const object of ordered) {
    if (object.parentId && !ids.has(object.parentId) && !result.some((e) => e.object === object))
      result.push({ object, depth: 0 });
  }
  return result;
}

export function findGroup(tree: ObjectTree, groupId: string): ObjectGroup | undefined {
  return tree.groups?.find((g) => g.id === groupId);
}

/** The top-level objects listed in the group, in tree order. */
export function groupMemberIds(tree: ObjectTree, groupId: string): string[] {
  return groupedOrder(tree.objects)
    .filter((o) => !o.parentId && o.groupId === groupId)
    .map((o) => o.id);
}

/** Only top-level objects go into a group – a child is listed under its parent. */
export function canGroup(tree: ObjectTree, ids: readonly string[]): boolean {
  return ids.length > 0 && ids.every((id) => findObject(tree, id)?.parentId === '');
}

/** The list id of a group's row: never an object id, which has no colon. */
export function groupRowId(groupId: string): string {
  return `group:${groupId}`;
}

export function groupIdOfRow(rowId: string): string | null {
  return rowId.startsWith('group:') ? rowId.slice('group:'.length) : null;
}

export type HierarchyRow =
  | { kind: 'object'; object: SceneObject; depth: number; hasChildren: boolean }
  | { kind: 'group'; group: ObjectGroup; memberIds: string[]; depth: number };

/**
 * Tree order with a row in front of each group and the subtrees of collapsed objects and groups
 * left out. `expanded` holds object ids and group row ids; anything not in it counts as collapsed,
 * so a freshly loaded scene shows only its top level.
 */
export function hierarchyRows(tree: ObjectTree, expanded: ReadonlySet<string>): HierarchyRow[] {
  const parents = new Set(tree.objects.map((o) => o.parentId).filter(Boolean));
  const groups = new Map((tree.groups ?? []).map((g) => [g.id, g]));
  const result: HierarchyRow[] = [];
  // Depth of the collapsed row whose subtree is being skipped (-1 = showing rows).
  let collapsedAt = -1;
  // The group of the top-level object the current entries belong to, and whether it is folded.
  let group: ObjectGroup | undefined;
  let groupFolded = false;
  for (const { object, depth } of treeOrder(tree)) {
    if (depth === 0) {
      const next = object.parentId ? undefined : groups.get(object.groupId ?? '');
      if (next && next !== group) {
        const rowId = groupRowId(next.id);
        result.push({ kind: 'group', group: next, memberIds: groupMemberIds(tree, next.id), depth: 0 });
        groupFolded = !expanded.has(rowId);
      }
      group = next;
    }
    if (group && groupFolded) continue;
    if (collapsedAt >= 0) {
      if (depth > collapsedAt) continue;
      collapsedAt = -1;
    }
    const hasChildren = parents.has(object.id);
    result.push({ kind: 'object', object, depth: depth + (group ? 1 : 0), hasChildren });
    if (hasChildren && !expanded.has(object.id)) collapsedAt = depth;
  }
  return result;
}