image-studio/commands/document-tree.ts

import type { ImageDocument } from "@core/document";
import type { ArtboardId, LayerId } from "@core/id";
import type { Layer } from "@core/layer";
import type { LayerGroup } from "@core/layer-group";
import { getLayerMask } from "@core/layer-mask-utils";
import type { MaskEditState } from "@editor/state";

export type LayerLocation = {
  artboardId: ArtboardId;
  parentGroupId?: LayerId;
  index: number;
  layer: Layer;
  siblings: readonly Layer[];
};

export function findLayerLocation(document: ImageDocument, layerId: LayerId): LayerLocation | undefined {
  for (const artboard of document.artboards) {
    const location = findLayerLocationInTree(artboard.layers, layerId, artboard.id);
    if (location) return location;
  }
  return undefined;
}

export function isReferencedMaskLayer(document: ImageDocument, maskLayerId: LayerId): boolean {
  return document.artboards.some((artboard) => isReferencedMaskLayerInTree(artboard.layers, maskLayerId));
}

export function isReferencedMaskLayerInTree(layers: readonly Layer[], maskLayerId: LayerId): boolean {
  for (const layer of layers) {
    if (getLayerMask(layer)?.maskLayerId === maskLayerId) return true;
    if (layer.type === "group" && isReferencedMaskLayerInTree(layer.children, maskLayerId)) return true;
  }
  return false;
}

export function findLayerLocationInTree(layers: Layer[], layerId: LayerId, artboardId: ArtboardId, parentGroupId?: LayerId): LayerLocation | undefined {
  for (let index = 0; index < layers.length; index++) {
    const layer = layers[index];
    if (!layer) continue;
    if (layer.id === layerId) return { artboardId, parentGroupId, index, layer, siblings: layers };
    if (layer.type === "group") {
      const child = findLayerLocationInTree(layer.children, layerId, artboardId, layer.id);
      if (child) return child;
    }
  }
  return undefined;
}

export function mapLayerInDocument(document: ImageDocument, layerId: LayerId, mapLayer: (layer: Layer) => Layer): ImageDocument {
  return {
    ...document,
    artboards: document.artboards.map((artboard) => ({ ...artboard, layers: mapLayerInTree(artboard.layers, layerId, mapLayer) })),
  };
}

export function mapLayerInTree(layers: Layer[], layerId: LayerId, mapLayer: (layer: Layer) => Layer): Layer[] {
  return layers.map((layer) => {
    if (layer.id === layerId) return mapLayer(layer);
    if (layer.type === "group") return { ...layer, children: mapLayerInTree(layer.children, layerId, mapLayer) };
    return layer;
  });
}

export function insertLayer(document: ImageDocument, artboardId: ArtboardId, parentGroupId: LayerId | undefined, layer: Layer, index?: number): ImageDocument {
  return {
    ...document,
    artboards: document.artboards.map((artboard) => {
      if (artboard.id !== artboardId) return artboard;
      if (!parentGroupId) return { ...artboard, layers: insertAt(artboard.layers, layer, index) };
      return { ...artboard, layers: insertLayerInGroup(artboard.layers, parentGroupId, layer, index) };
    }),
  };
}

export function insertLayerInGroup(layers: Layer[], groupId: LayerId, layer: Layer, index?: number): Layer[] {
  return layers.map((candidate) => {
    if (candidate.type === "group" && candidate.id === groupId) return { ...candidate, children: insertAt(candidate.children, layer, index) };
    if (candidate.type === "group") return { ...candidate, children: insertLayerInGroup(candidate.children, groupId, layer, index) };
    return candidate;
  });
}

export function replaceLayerListInDocument(document: ImageDocument, artboardId: ArtboardId, parentGroupId: LayerId | undefined, layers: Layer[]): ImageDocument {
  return {
    ...document,
    artboards: document.artboards.map((artboard) => {
      if (artboard.id !== artboardId) return artboard;
      if (!parentGroupId) return { ...artboard, layers };
      return { ...artboard, layers: replaceLayerListInGroup(artboard.layers, parentGroupId, layers) };
    }),
  };
}

export function replaceLayerListInGroup(layers: Layer[], groupId: LayerId, children: Layer[]): Layer[] {
  return layers.map((layer) => {
    if (layer.type === "group" && layer.id === groupId) return { ...layer, children };
    if (layer.type === "group") return { ...layer, children: replaceLayerListInGroup(layer.children, groupId, children) };
    return layer;
  });
}

export function replaceSelectedLayersWithGroup(layers: readonly Layer[], selectedLayerIds: ReadonlySet<LayerId>, group: LayerGroup): Layer[] {
  const next: Layer[] = [];
  let inserted = false;
  for (const layer of layers) {
    if (!selectedLayerIds.has(layer.id)) {
      next.push(layer);
      continue;
    }
    if (!inserted) {
      next.push(group);
      inserted = true;
    }
  }
  return next;
}

export function removeLayerFromDocument(document: ImageDocument, layerId: LayerId): { document: ImageDocument; layer?: Layer } {
  let removed: Layer | undefined;
  return {
    document: {
      ...document,
      artboards: document.artboards.map((artboard) => {
        const result = removeLayerFromTree(artboard.layers, layerId);
        if (result.layer) removed = result.layer;
        return { ...artboard, layers: result.layers };
      }),
    },
    layer: removed,
  };
}

export function removeLayerFromTree(layers: Layer[], layerId: LayerId): { layers: Layer[]; layer?: Layer } {
  let removed: Layer | undefined;
  const next: Layer[] = [];
  for (const layer of layers) {
    if (layer.id === layerId) {
      removed = layer;
      continue;
    }
    if (layer.type === "group") {
      const result = removeLayerFromTree(layer.children, layerId);
      if (result.layer) removed = result.layer;
      next.push({ ...layer, children: result.layers });
    } else {
      next.push(layer);
    }
  }
  return { layers: next, layer: removed };
}

export function ungroupLayerInDocument(document: ImageDocument, groupId: LayerId): { document: ImageDocument; changed: boolean; artboardId?: ArtboardId; children: Layer[] } {
  let changed = false;
  let artboardId: ArtboardId | undefined;
  let children: Layer[] = [];
  const next = {
    ...document,
    artboards: document.artboards.map((artboard) => {
      const result = ungroupLayerInTree(artboard.layers, groupId);
      if (result.changed) {
        changed = true;
        artboardId = artboard.id;
        children = result.children;
      }
      return { ...artboard, layers: result.layers };
    }),
  };
  return { document: next, changed, artboardId, children };
}

export function ungroupLayerInTree(layers: Layer[], groupId: LayerId): { layers: Layer[]; changed: boolean; children: Layer[] } {
  const next: Layer[] = [];
  for (const layer of layers) {
    if (layer.type === "group" && layer.id === groupId) return { layers: [...next, ...layer.children, ...layers.slice(next.length + 1)], changed: true, children: layer.children };
    if (layer.type === "group") {
      const result = ungroupLayerInTree(layer.children, groupId);
      if (result.changed) return { layers: [...next, { ...layer, children: result.layers }, ...layers.slice(next.length + 1)], changed: true, children: result.children };
    }
    next.push(layer);
  }
  return { layers, changed: false, children: [] };
}

export function insertAt(layers: Layer[], layer: Layer, index = layers.length) {
  const clamped = Math.max(0, Math.min(index, layers.length));
  return [...layers.slice(0, clamped), layer, ...layers.slice(clamped)];
}

export function findGroup(document: ImageDocument, groupId: LayerId): LayerGroup | undefined {
  for (const artboard of document.artboards) {
    const group = findGroupInTree(artboard.layers, groupId);
    if (group) return group;
  }
  return undefined;
}

export function findGroupInTree(layers: Layer[], groupId: LayerId): LayerGroup | undefined {
  for (const layer of layers) {
    if (layer.type === "group" && layer.id === groupId) return layer;
    if (layer.type === "group") {
      const child = findGroupInTree(layer.children, groupId);
      if (child) return child;
    }
  }
  return undefined;
}

export function removeLayerMaskReference(layer: Layer): Layer {
  const next = { ...layer };
  delete next.layerMask;
  delete next.clippingMask;
  return next;
}

export function withLayerMask(layer: Layer, maskLayerId: LayerId): Layer {
  return {
    ...layer,
    layerMask: {
      kind: "raster",
      maskLayerId,
      enabled: true,
      inverted: false,
    },
    clippingMask: { maskLayerId },
  };
}

export function removeUnreferencedMaskLayer(document: ImageDocument, maskLayerId: LayerId): ImageDocument {
  if (isMaskLayerReferenced(document, maskLayerId)) return document;
  return removeLayerFromDocument(document, maskLayerId).document;
}

export function isMaskLayerReferenced(document: ImageDocument, maskLayerId: LayerId): boolean {
  return collectClippingMaskIds(document.artboards.flatMap((artboard) => artboard.layers)).has(maskLayerId);
}

export function removeMissingMaskReferences(document: ImageDocument): ImageDocument {
  const existingLayerIds = collectDocumentLayerIds(document);
  return mapAllLayersInDocument(document, (layer) => {
    const mask = getLayerMask(layer);
    if (!mask || existingLayerIds.has(mask.maskLayerId)) return layer;
    return removeLayerMaskReference(layer);
  });
}

export function isMaskEditFor(maskEdit: MaskEditState | undefined, targetLayerId: LayerId, maskLayerId: LayerId) {
  return maskEdit?.targetLayerId === targetLayerId && maskEdit.maskLayerId === maskLayerId;
}

export function isMaskEditValid(maskEdit: MaskEditState | undefined, document: ImageDocument) {
  if (!maskEdit) return false;
  const target = findLayerLocation(document, maskEdit.targetLayerId)?.layer;
  if (maskEdit.kind === "inpaintRegion") {
    const region = document.inpaintRegions.find((candidate) => candidate.id === maskEdit.inpaintRegionId && candidate.targetLayerId === maskEdit.targetLayerId && candidate.maskAssetId === maskEdit.maskAssetId);
    return Boolean(target && region && document.assets.some((asset) => asset.id === region.maskAssetId));
  }
  if (!maskEdit.maskLayerId) return false;
  const mask = findLayerLocation(document, maskEdit.maskLayerId)?.layer;
  return Boolean(target && getLayerMask(target)?.maskLayerId === maskEdit.maskLayerId && mask && mask.type !== "group");
}

export function collectDocumentLayerIds(document: ImageDocument): Set<LayerId> {
  const ids = new Set<LayerId>();
  for (const artboard of document.artboards) collectLayerIdsFromTree(artboard.layers, ids);
  return ids;
}

export function collectLayerIds(layer: Layer, ids = new Set<LayerId>()): Set<LayerId> {
  ids.add(layer.id);
  if (layer.type === "group") collectLayerIdsFromTree(layer.children, ids);
  return ids;
}

export function collectLayerIdsFromTree(layers: readonly Layer[], ids = new Set<LayerId>()): Set<LayerId> {
  for (const layer of layers) collectLayerIds(layer, ids);
  return ids;
}

export function collectClippingMaskIds(layers: readonly Layer[], ids = new Set<LayerId>()): Set<LayerId> {
  for (const layer of layers) {
    const mask = getLayerMask(layer);
    if (mask) ids.add(mask.maskLayerId);
    if (layer.type === "group") collectClippingMaskIds(layer.children, ids);
  }
  return ids;
}

export function collectAttachedMaskIds(layers: readonly Layer[], layerIds: readonly LayerId[]): LayerId[] {
  const layerIdSet = new Set(layerIds);
  return layers.flatMap((layer) => {
    const mask = getLayerMask(layer);
    return layerIdSet.has(layer.id) && mask ? [mask.maskLayerId] : [];
  });
}

export function mapAllLayersInDocument(document: ImageDocument, mapLayer: (layer: Layer) => Layer): ImageDocument {
  return {
    ...document,
    artboards: document.artboards.map((artboard) => ({ ...artboard, layers: mapAllLayersInTree(artboard.layers, mapLayer) })),
  };
}

export function mapAllLayersInTree(layers: Layer[], mapLayer: (layer: Layer) => Layer): Layer[] {
  return layers.map((layer) => {
    const mapped = layer.type === "group" ? { ...layer, children: mapAllLayersInTree(layer.children, mapLayer) } : layer;
    return mapLayer(mapped);
  });
}

export function removeInpaintRegionsForTargets(document: ImageDocument, targetLayerIds: ReadonlySet<LayerId>): ImageDocument {
  const removed = document.inpaintRegions.filter((region) => targetLayerIds.has(region.targetLayerId));
  if (removed.length === 0) return document;
  const inpaintRegions = document.inpaintRegions.filter((region) => !targetLayerIds.has(region.targetLayerId));
  const removedMaskAssetIds = new Set(removed.map((region) => region.maskAssetId));
  const retainedRegionAssetIds = new Set(inpaintRegions.map((region) => region.maskAssetId));
  const layerAssetIds = new Set<string>();
  for (const artboard of document.artboards) {
    const stack = [...artboard.layers];
    while (stack.length > 0) {
      const layer = stack.pop();
      if (!layer) continue;
      if (layer.type === "group") stack.push(...layer.children);
      else if (layer.type === "image" || layer.type === "raster") layerAssetIds.add(layer.assetId);
    }
  }
  return {
    ...document,
    inpaintRegions,
    assets: document.assets.filter((asset) => !removedMaskAssetIds.has(asset.id) || retainedRegionAssetIds.has(asset.id) || layerAssetIds.has(asset.id)),
  };
}