image-studio/commands/tool.ts

import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
import type { LayerId, ArtboardId, AssetId } from "@core/id";
import type { Layer } from "@core/layer";
import type { MaskViewMode } from "@editor/state";
import type { BrushSettings, ChromaKeySettings, GenerateSettings, MagicWandSettings, ToolId } from "@editor/tools";
import type { Command } from "./command";
import { commandIds } from "./ids";

export type ToolSetActivePayload = {
  tool: ToolId;
};

export type ToolSetBrushSettingsPayload = Partial<BrushSettings>;

export type ToolSetGenerateSettingsPayload = Partial<GenerateSettings>;

export type ToolSetChromaKeySettingsPayload = Partial<ChromaKeySettings>;

export type ToolSetMagicWandSettingsPayload = Partial<MagicWandSettings>;

export type ToolSetBrushPreviewPayload = { position: Vec2D } | undefined;

export type ToolSetBrushStrokePreviewPayload =
  | {
      layerId: LayerId;
      assetId: AssetId;
      source: string;
    }
  | undefined;

export type ToolSetMaskViewModePayload = {
  mode: MaskViewMode;
};

export type ToolEnterMaskEditPayload = {
  targetLayerId: LayerId;
  maskLayerId: LayerId;
};

export const toolSetActiveCommand: Command<ToolSetActivePayload> = {
  id: commandIds.toolSetActive,
  name: "Set active tool",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          activeTool: payload.tool,
          interactionMode: { type: "tool", tool: payload.tool },
        },
        brushPreview: undefined,
        brushStrokePreview: undefined,
      },
    };
  },
};

export const toolSetGenerateSettingsCommand: Command<ToolSetGenerateSettingsPayload> = {
  id: commandIds.toolSetGenerateSettings,
  name: "Set generate settings",
  execute({ state }, payload) {
    const mode = payload.mode ?? state.editor.tools.generate.mode;
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          generate: {
            mode,
            model: payload.model ?? state.editor.tools.generate.model,
            prompt: payload.prompt ?? state.editor.tools.generate.prompt,
            negativePrompt: payload.negativePrompt ?? state.editor.tools.generate.negativePrompt,
            strength: clampNumber(payload.strength ?? state.editor.tools.generate.strength, 0, 100),
            steps: Math.round(clampNumber(payload.steps ?? state.editor.tools.generate.steps, 1, 150)),
            cfg: clampNumber(payload.cfg ?? state.editor.tools.generate.cfg, 0, 30),
            seed: Math.round(clampNumber(payload.seed ?? state.editor.tools.generate.seed, -1, Number.MAX_SAFE_INTEGER)),
            sampler: payload.sampler ?? state.editor.tools.generate.sampler,
            scheduler: payload.scheduler ?? state.editor.tools.generate.scheduler,
            width: Math.round(clampNumber(payload.width ?? state.editor.tools.generate.width, 64, 4096)),
            height: Math.round(clampNumber(payload.height ?? state.editor.tools.generate.height, 64, 4096)),
            outpaint: {
              left: Math.round(clampNumber(payload.outpaint?.left ?? state.editor.tools.generate.outpaint.left, 0, 2048)),
              top: Math.round(clampNumber(payload.outpaint?.top ?? state.editor.tools.generate.outpaint.top, 0, 2048)),
              right: Math.round(clampNumber(payload.outpaint?.right ?? state.editor.tools.generate.outpaint.right, 0, 2048)),
              bottom: Math.round(clampNumber(payload.outpaint?.bottom ?? state.editor.tools.generate.outpaint.bottom, 0, 2048)),
              feathering: Math.round(clampNumber(payload.outpaint?.feathering ?? state.editor.tools.generate.outpaint.feathering, 0, 512)),
            },
          },
        },
      },
    };
  },
};

export const toolSetBrushSettingsCommand: Command<ToolSetBrushSettingsPayload> = {
  id: commandIds.toolSetBrushSettings,
  name: "Set brush settings",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          brush: {
            color: payload.color ?? state.editor.tools.brush.color,
            size: clampNumber(payload.size ?? state.editor.tools.brush.size, 1, 200),
            hardness: clampNumber(payload.hardness ?? state.editor.tools.brush.hardness, 0, 100),
          },
        },
      },
    };
  },
};

export const toolSetChromaKeySettingsCommand: Command<ToolSetChromaKeySettingsPayload> = {
  id: commandIds.toolSetChromaKeySettings,
  name: "Set chroma key settings",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          chromaKey: {
            color: payload.color ?? state.editor.tools.chromaKey.color,
            tolerance: clampNumber(payload.tolerance ?? state.editor.tools.chromaKey.tolerance, 0, 255),
            softness: clampNumber(payload.softness ?? state.editor.tools.chromaKey.softness, 0, 255),
            feather: clampNumber(payload.feather ?? state.editor.tools.chromaKey.feather, 0, 20),
            choke: clampNumber(payload.choke ?? state.editor.tools.chromaKey.choke, -20, 20),
            despeckle: clampNumber(payload.despeckle ?? state.editor.tools.chromaKey.despeckle, 0, 20),
            spill: clampNumber(payload.spill ?? state.editor.tools.chromaKey.spill, 0, 100),
          },
        },
      },
    };
  },
};

export const toolSetMagicWandSettingsCommand: Command<ToolSetMagicWandSettingsPayload> = {
  id: commandIds.toolSetMagicWandSettings,
  name: "Set magic wand settings",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          magicWand: {
            tolerance: clampNumber(payload.tolerance ?? state.editor.tools.magicWand.tolerance, 0, 255),
            feather: clampNumber(payload.feather ?? state.editor.tools.magicWand.feather, 0, 20),
            choke: clampNumber(payload.choke ?? state.editor.tools.magicWand.choke, -20, 20),
            despeckle: clampNumber(payload.despeckle ?? state.editor.tools.magicWand.despeckle, 0, 20),
            contiguous: payload.contiguous ?? state.editor.tools.magicWand.contiguous,
            mode: payload.mode ?? state.editor.tools.magicWand.mode,
          },
        },
      },
    };
  },
};

export const toolSetBrushPreviewCommand: Command<ToolSetBrushPreviewPayload> = {
  id: commandIds.toolSetBrushPreview,
  name: "Set brush preview",
  execute({ state }, payload) {
    if (!payload) {
      if (!state.editor.brushPreview) return state;
      return { ...state, editor: { ...state.editor, brushPreview: undefined } };
    }

    const currentPosition = state.editor.brushPreview?.position;
    if (currentPosition && currentPosition.x === payload.position.x && currentPosition.y === payload.position.y) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        brushPreview: { position: { ...payload.position } },
      },
    };
  },
};

export const toolSetBrushStrokePreviewCommand: Command<ToolSetBrushStrokePreviewPayload> = {
  id: commandIds.toolSetBrushStrokePreview,
  name: "Set brush stroke preview",
  execute({ state }, payload) {
    if (!payload) {
      if (!state.editor.brushStrokePreview) return state;
      return { ...state, editor: { ...state.editor, brushStrokePreview: undefined } };
    }

    const currentPreview = state.editor.brushStrokePreview;
    if (currentPreview?.layerId === payload.layerId && currentPreview.assetId === payload.assetId && currentPreview.source === payload.source) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        brushStrokePreview: { layerId: payload.layerId, assetId: payload.assetId, source: payload.source },
      },
    };
  },
};

export const toolSetMaskViewModeCommand: Command<ToolSetMaskViewModePayload> = {
  id: commandIds.toolSetMaskViewMode,
  name: "Set mask view mode",
  execute({ state }, payload) {
    if (!state.editor.maskEdit || state.editor.maskEdit.viewMode === payload.mode) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        maskEdit: { ...state.editor.maskEdit, viewMode: payload.mode },
      },
    };
  },
};

export const toolEnterMaskEditCommand: Command<ToolEnterMaskEditPayload> = {
  id: commandIds.toolEnterMaskEdit,
  name: "Enter mask edit",
  execute({ state }, payload) {
    const targetLocation = findLayerLocation(state.document, payload.targetLayerId);
    const maskLocation = findLayerLocation(state.document, payload.maskLayerId);
    if (!targetLocation || !maskLocation) return state;
    if (targetLocation.layer.clippingMask?.maskLayerId !== payload.maskLayerId) return state;
    if (maskLocation.layer.type === "group") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        selection: { artboardId: targetLocation.artboardId, layerIds: [payload.targetLayerId] },
        maskEdit: { targetLayerId: payload.targetLayerId, maskLayerId: payload.maskLayerId },
        brushPreview: undefined,
        brushStrokePreview: undefined,
        tools: {
          ...state.editor.tools,
          activeTool: "brush",
          interactionMode: { type: "tool", tool: "brush" },
        },
      },
    };
  },
};

export const toolExitMaskEditCommand: Command = {
  id: commandIds.toolExitMaskEdit,
  name: "Exit mask edit",
  execute({ state }) {
    if (!state.editor.maskEdit) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        maskEdit: undefined,
        brushPreview: undefined,
        brushStrokePreview: undefined,
      },
    };
  },
};

export const toolEnterTemporaryPanCommand: Command = {
  id: commandIds.toolEnterTemporaryPan,
  name: "Enter temporary pan",
  execute({ state }) {
    if (state.editor.tools.interactionMode.type === "temporary-pan") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          interactionMode: { type: "temporary-pan", previousTool: state.editor.tools.activeTool },
        },
      },
    };
  },
};

export const toolExitTemporaryPanCommand: Command = {
  id: commandIds.toolExitTemporaryPan,
  name: "Exit temporary pan",
  execute({ state }) {
    const mode = state.editor.tools.interactionMode;
    if (mode.type !== "temporary-pan") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          activeTool: mode.previousTool,
          interactionMode: { type: "tool", tool: mode.previousTool },
        },
      },
    };
  },
};

export const toolCommands = [
  toolSetActiveCommand,
  toolSetGenerateSettingsCommand,
  toolSetBrushSettingsCommand,
  toolSetChromaKeySettingsCommand,
  toolSetMagicWandSettingsCommand,
  toolSetBrushPreviewCommand,
  toolSetBrushStrokePreviewCommand,
  toolSetMaskViewModeCommand,
  toolEnterMaskEditCommand,
  toolExitMaskEditCommand,
  toolEnterTemporaryPanCommand,
  toolExitTemporaryPanCommand,
] satisfies Command<unknown>[];

function clampNumber(value: number, min: number, max: number) {
  if (!Number.isFinite(value)) return min;
  return Math.max(min, Math.min(max, value));
}

type LayerLocation = {
  artboardId: ArtboardId;
  layer: Layer;
};

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

function findLayerInTree(layers: readonly Layer[], layerId: LayerId): Layer | undefined {
  for (const layer of layers) {
    if (layer.id === layerId) return layer;
    if (layer.type === "group") {
      const child = findLayerInTree(layer.children, layerId);
      if (child) return child;
    }
  }
  return undefined;
}