image-studio/commands/document.ts

import type { Asset } from "@core/asset";
import type { ImageDocument } from "@core/document";
import type { Rect } from "@core/geometry";
import type { ArtboardId, LayerId } from "@core/id";
import type { ImageLayer } from "@core/image-layer";
import type { Layer } from "@core/layer";
import type { LayerGroup } from "@core/layer-group";
import type { Command } from "./command";
import { commandIds } from "./ids";

export type DocumentAddArtboardPayload = {
  id: ArtboardId;
  name: string;
  bounds: Rect;
};

export type DocumentSetArtboardBoundsPayload = {
  id: ArtboardId;
  bounds: Rect;
};

export type DocumentRemoveArtboardPayload = {
  id: ArtboardId;
};

export type DocumentSetArtboardVisiblePayload = {
  id: ArtboardId;
  visible: boolean;
};

export type DocumentSetArtboardLockedPayload = {
  id: ArtboardId;
  locked: boolean;
};

export type DocumentRenameArtboardPayload = {
  id: ArtboardId;
  name: string;
};

export type DocumentAddAssetPayload = {
  asset: Asset;
};

export type DocumentAddImageLayerPayload = {
  artboardId: ArtboardId;
  parentGroupId?: LayerId;
  layer: ImageLayer;
};

export type DocumentAddGroupLayerPayload = {
  artboardId: ArtboardId;
  parentGroupId?: LayerId;
  group: LayerGroup;
};

export type DocumentMoveLayerPayload = {
  layerId: LayerId;
  toArtboardId: ArtboardId;
  toParentGroupId?: LayerId;
  toIndex: number;
};

export type DocumentGroupLayersPayload = {
  artboardId: ArtboardId;
  layerIds: LayerId[];
  group: LayerGroup;
};

export type DocumentUngroupLayerPayload = {
  groupId: LayerId;
};

export type DocumentRemoveLayerPayload = {
  layerId: LayerId;
};

export type DocumentSetLayerVisiblePayload = {
  layerId: LayerId;
  visible: boolean;
};

export type DocumentSetLayerLockedPayload = {
  layerId: LayerId;
  locked: boolean;
};

export type DocumentRenameLayerPayload = {
  layerId: LayerId;
  name: string;
};

export type DocumentSetLayerClippingMaskPayload = {
  layerId: LayerId;
  maskLayerId?: LayerId;
};

export const documentAddArtboardCommand: Command<DocumentAddArtboardPayload> = {
  id: commandIds.documentAddArtboard,
  name: "Add artboard",
  execute({ state }, payload) {
    return {
      ...state,
      document: {
        ...state.document,
        artboards: [
          ...state.document.artboards,
          {
            id: payload.id,
            name: payload.name,
            bounds: payload.bounds,
            backgroundColor: "transparent",
            visible: true,
            locked: false,
            layers: [],
          },
        ],
      },
    };
  },
};

export const documentSetArtboardBoundsCommand: Command<DocumentSetArtboardBoundsPayload> = {
  id: commandIds.documentSetArtboardBounds,
  name: "Set artboard bounds",
  execute({ state }, payload) {
    return {
      ...state,
      document: {
        ...state.document,
        artboards: state.document.artboards.map((artboard) =>
          artboard.id === payload.id ? { ...artboard, bounds: { ...payload.bounds } } : artboard,
        ),
      },
    };
  },
};

export const documentRemoveArtboardCommand: Command<DocumentRemoveArtboardPayload> = {
  id: commandIds.documentRemoveArtboard,
  name: "Remove artboard",
  execute({ state }, payload) {
    const removedSelectedArtboard = state.editor.selection.artboardId === payload.id;
    return {
      ...state,
      document: {
        ...state.document,
        artboards: state.document.artboards.filter((artboard) => artboard.id !== payload.id),
      },
      editor: removedSelectedArtboard ? { ...state.editor, selection: { layerIds: [] } } : state.editor,
    };
  },
};

export const documentSetArtboardVisibleCommand: Command<DocumentSetArtboardVisiblePayload> = {
  id: commandIds.documentSetArtboardVisible,
  name: "Set artboard visible",
  execute({ state }, payload) {
    return {
      ...state,
      document: {
        ...state.document,
        artboards: state.document.artboards.map((artboard) => (artboard.id === payload.id ? { ...artboard, visible: payload.visible } : artboard)),
      },
    };
  },
};

export const documentSetArtboardLockedCommand: Command<DocumentSetArtboardLockedPayload> = {
  id: commandIds.documentSetArtboardLocked,
  name: "Set artboard locked",
  execute({ state }, payload) {
    return {
      ...state,
      document: {
        ...state.document,
        artboards: state.document.artboards.map((artboard) => (artboard.id === payload.id ? { ...artboard, locked: payload.locked } : artboard)),
      },
    };
  },
};

export const documentRenameArtboardCommand: Command<DocumentRenameArtboardPayload> = {
  id: commandIds.documentRenameArtboard,
  name: "Rename artboard",
  execute({ state }, payload) {
    const name = payload.name.trim();
    if (!name) return state;
    return {
      ...state,
      document: {
        ...state.document,
        artboards: state.document.artboards.map((artboard) => (artboard.id === payload.id ? { ...artboard, name } : artboard)),
      },
    };
  },
};

export const documentAddAssetCommand: Command<DocumentAddAssetPayload> = {
  id: commandIds.documentAddAsset,
  name: "Add asset",
  execute({ state }, payload) {
    if (state.document.assets.some((asset) => asset.id === payload.asset.id)) return state;

    return {
      ...state,
      document: {
        ...state.document,
        assets: [...state.document.assets, payload.asset],
      },
    };
  },
};

export const documentAddImageLayerCommand: Command<DocumentAddImageLayerPayload> = {
  id: commandIds.documentAddImageLayer,
  name: "Add image layer",
  execute({ state }, payload) {
    return {
      ...state,
      document: insertLayer(state.document, payload.artboardId, payload.parentGroupId, payload.layer),
    };
  },
};

export const documentAddGroupLayerCommand: Command<DocumentAddGroupLayerPayload> = {
  id: commandIds.documentAddGroupLayer,
  name: "Add group layer",
  execute({ state }, payload) {
    return {
      ...state,
      document: insertLayer(state.document, payload.artboardId, payload.parentGroupId, payload.group),
      editor: { ...state.editor, selection: { artboardId: payload.artboardId, layerIds: [payload.group.id] } },
    };
  },
};

export const documentMoveLayerCommand: Command<DocumentMoveLayerPayload> = {
  id: commandIds.documentMoveLayer,
  name: "Move layer",
  execute({ state }, payload) {
    if (payload.toParentGroupId && !findGroup(state.document, payload.toParentGroupId)) return state;

    const removed = removeLayerFromDocument(state.document, payload.layerId);
    if (!removed.layer) return state;
    if (payload.toParentGroupId && !findGroup(removed.document, payload.toParentGroupId)) return state;

    return {
      ...state,
      document: insertLayer(removed.document, payload.toArtboardId, payload.toParentGroupId, removed.layer, payload.toIndex),
    };
  },
};

export const documentGroupLayersCommand: Command<DocumentGroupLayersPayload> = {
  id: commandIds.documentGroupLayers,
  name: "Group layers",
  execute({ state }, payload) {
    const uniqueIds = [...new Set(payload.layerIds)];
    if (uniqueIds.length === 0) return state;

    const artboard = state.document.artboards.find((candidate) => candidate.id === payload.artboardId);
    if (!artboard) return state;

    const selected = artboard.layers.filter((layer) => uniqueIds.includes(layer.id));
    if (selected.length === 0) return state;

    const firstIndex = artboard.layers.findIndex((layer) => layer.id === selected[0]?.id);
    const group: LayerGroup = { ...payload.group, children: selected };
    const document = {
      ...state.document,
      artboards: state.document.artboards.map((candidate) =>
        candidate.id === payload.artboardId
          ? { ...candidate, layers: [...candidate.layers.filter((layer) => !uniqueIds.includes(layer.id)).slice(0, firstIndex), group, ...candidate.layers.filter((layer) => !uniqueIds.includes(layer.id)).slice(firstIndex)] }
          : candidate,
      ),
    };

    return {
      ...state,
      document,
      editor: { ...state.editor, selection: { artboardId: payload.artboardId, layerIds: [group.id] } },
    };
  },
};

export const documentUngroupLayerCommand: Command<DocumentUngroupLayerPayload> = {
  id: commandIds.documentUngroupLayer,
  name: "Ungroup layer",
  execute({ state }, payload) {
    const result = ungroupLayerInDocument(state.document, payload.groupId);
    if (!result.changed) return state;

    return {
      ...state,
      document: result.document,
      editor: { ...state.editor, selection: { artboardId: result.artboardId, layerIds: result.children.map((layer) => layer.id) } },
    };
  },
};

export const documentSetLayerVisibleCommand: Command<DocumentSetLayerVisiblePayload> = {
  id: commandIds.documentSetLayerVisible,
  name: "Set layer visible",
  execute({ state }, payload) {
    return { ...state, document: mapLayerInDocument(state.document, payload.layerId, (layer) => ({ ...layer, visible: payload.visible })) };
  },
};

export const documentSetLayerLockedCommand: Command<DocumentSetLayerLockedPayload> = {
  id: commandIds.documentSetLayerLocked,
  name: "Set layer locked",
  execute({ state }, payload) {
    return { ...state, document: mapLayerInDocument(state.document, payload.layerId, (layer) => ({ ...layer, locked: payload.locked })) };
  },
};

export const documentRenameLayerCommand: Command<DocumentRenameLayerPayload> = {
  id: commandIds.documentRenameLayer,
  name: "Rename layer",
  execute({ state }, payload) {
    const name = payload.name.trim();
    if (!name) return state;
    return { ...state, document: mapLayerInDocument(state.document, payload.layerId, (layer) => ({ ...layer, name })) };
  },
};

export const documentSetLayerClippingMaskCommand: Command<DocumentSetLayerClippingMaskPayload> = {
  id: commandIds.documentSetLayerClippingMask,
  name: "Set layer clipping mask",
  execute({ state }, payload) {
    if (payload.maskLayerId === payload.layerId) return state;

    if (!payload.maskLayerId) {
      return {
        ...state,
        document: mapLayerInDocument(state.document, payload.layerId, (layer) => {
          const { clippingMask: _clippingMask, ...rest } = layer;
          return rest;
        }),
      };
    }

    const targetLocation = findLayerLocation(state.document, payload.layerId);
    const maskLocation = findLayerLocation(state.document, payload.maskLayerId);
    if (!targetLocation || !maskLocation) return state;
    if (targetLocation.artboardId !== maskLocation.artboardId || targetLocation.parentGroupId !== maskLocation.parentGroupId) return state;

    const removed = removeLayerFromDocument(state.document, payload.layerId);
    if (!removed.layer) return state;

    const maskLocationAfterRemoval = findLayerLocation(removed.document, payload.maskLayerId);
    if (!maskLocationAfterRemoval) return state;

    return {
      ...state,
      document: insertLayer(
        removed.document,
        maskLocationAfterRemoval.artboardId,
        maskLocationAfterRemoval.parentGroupId,
        { ...removed.layer, clippingMask: { maskLayerId: payload.maskLayerId } },
        maskLocationAfterRemoval.index + 1,
      ),
    };
  },
};

export const documentRemoveLayerCommand: Command<DocumentRemoveLayerPayload> = {
  id: commandIds.documentRemoveLayer,
  name: "Remove layer",
  execute({ state }, payload) {
    const removed = removeLayerFromDocument(state.document, payload.layerId);
    if (!removed.layer) return state;

    return {
      ...state,
      document: removed.document,
      editor: { ...state.editor, selection: { ...state.editor.selection, layerIds: state.editor.selection.layerIds.filter((id) => id !== payload.layerId) } },
    };
  },
};

export const documentCommands = [
  documentAddArtboardCommand,
  documentSetArtboardBoundsCommand,
  documentRemoveArtboardCommand,
  documentSetArtboardVisibleCommand,
  documentSetArtboardLockedCommand,
  documentRenameArtboardCommand,
  documentAddAssetCommand,
  documentAddImageLayerCommand,
  documentAddGroupLayerCommand,
  documentMoveLayerCommand,
  documentGroupLayersCommand,
  documentUngroupLayerCommand,
  documentRemoveLayerCommand,
  documentSetLayerVisibleCommand,
  documentSetLayerLockedCommand,
  documentRenameLayerCommand,
  documentSetLayerClippingMaskCommand,
] satisfies Command<unknown>[];

type LayerLocation = {
  artboardId: ArtboardId;
  parentGroupId?: LayerId;
  index: number;
  layer: Layer;
};

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;
}

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 };
    if (layer.type === "group") {
      const child = findLayerLocationInTree(layer.children, layerId, artboardId, layer.id);
      if (child) return child;
    }
  }
  return undefined;
}

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) })),
  };
}

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;
  });
}

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) };
    }),
  };
}

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;
  });
}

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,
  };
}

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 };
}

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 };
}

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: [] };
}

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)];
}

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;
}

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;
}