image-studio/operations/paint/brush.ts

import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
import type { AssetId, LayerId } from "@core/id";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import type { RasterLayer } from "@core/raster-layer";
import type { MaskEditState, SelectionState } from "@editor/state";
import { isPanInteractionMode, type ToolState } from "@editor/tools";
import type { AppStore } from "@editor/store";
import { brushSurfaceDataUrl, brushSurfaceObjectUrl, cancelFrame, createBrushSurface, drawBrushSegment, drawFeatherSegment, releaseObjectUrl, scheduleFrame, type BrushSurface } from "@platform/browser/brushRaster";

export type BrushSession = {
  layerId: string;
  assetId: string;
  width: number;
  height: number;
  surface: BrushSurface;
  ready: Promise<boolean>;
  previousPoint: Vec2D;
  mode: "brush" | "eraser" | "feather";
  changed?: boolean;
  pending?: Promise<void>;
  cancelled?: boolean;
  previewClosed?: boolean;
  previewRequested?: boolean;
  previewInFlight?: boolean;
  previewFrame?: number;
  previewSource?: string;
  targetLayer: RasterLayer;
  featherRadius?: number;
  pendingLayerMask?: {
    targetLayerId: LayerId;
    asset: ImageDocument["assets"][number];
    maskLayer: RasterLayer;
  };
};

export type BrushTargetEditorState = {
  selection: SelectionState;
  tools: Pick<ToolState, "activeTool" | "interactionMode">;
  maskEdit?: MaskEditState;
};

export function beginBrushSession(document: ImageDocument, editor: BrushTargetEditorState, point: Vec2D): BrushSession | undefined {
  const target = resolveBrushTarget(document, editor);
  if (!target || !isPaintTool(editor.tools.activeTool)) return undefined;
  const { layer, asset } = target;

  const surface = createBrushSurface(asset.intrinsicSize.w, asset.intrinsicSize.h, asset.source);
  if (!surface) return undefined;

  const session: BrushSession = {
    layerId: layer.id,
    assetId: layer.assetId,
    width: surface.width,
    height: surface.height,
    surface,
    ready: surface.ready,
    previousPoint: point,
    mode: editor.tools.activeTool,
    targetLayer: layer,
    pendingLayerMask: target.pendingLayerMask,
  };
  return session;
}

export function canPreviewBrush(document: ImageDocument, editor: BrushTargetEditorState): boolean {
  if (editor.tools.activeTool === "feather" && !editor.maskEdit) return Boolean(resolveSelectedFeatherLayer(document, editor));
  return Boolean(resolveBrushTarget(document, editor));
}

export function brushUnavailableHint(document: ImageDocument, editor: BrushTargetEditorState): string | undefined {
  if (isPanInteractionMode(editor.tools.interactionMode) || !isPaintTool(editor.tools.activeTool)) return undefined;
  if (canPreviewBrush(document, editor)) return undefined;

  const layerId = editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0];
  if (!layerId) {
    if (editor.selection.artboardId) return "Brushes paint layers, not artboards. Select or add a raster layer first.";
    return "Select a raster layer or layer mask to paint.";
  }

  const layer = findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId);
  if (!layer) return "Select a raster layer or layer mask to paint.";
  if (layer.locked) return "Unlock this layer before painting.";
  if (layer.type === "image" && editor.tools.activeTool !== "feather") return "Image layers are non-destructive. Add a layer mask to paint or erase.";
  if (layer.type === "group") return "Select a raster layer inside the group to paint.";
  if (!editor.maskEdit && !layer.visible) return "Show this layer before painting.";
  return "Select a raster layer or layer mask to paint.";
}

function resolveBrushTarget(document: ImageDocument, editor: BrushTargetEditorState): { layer: RasterLayer; asset: ImageDocument["assets"][number]; pendingLayerMask?: BrushSession["pendingLayerMask"] } | undefined {
  if (isPanInteractionMode(editor.tools.interactionMode) || !isPaintTool(editor.tools.activeTool)) return undefined;
  if (editor.maskEdit?.kind === "inpaintRegion") {
    const target = findLayer(document.artboards.flatMap((artboard) => artboard.layers), editor.maskEdit.targetLayerId);
    const asset = document.assets.find((candidate) => candidate.id === editor.maskEdit?.maskAssetId);
    if (!target || (target.type !== "image" && target.type !== "raster") || target.locked || !asset) return undefined;
    return {
      layer: { ...target, type: "raster", assetId: asset.id },
      asset,
    };
  }
  const editingMask = Boolean(editor.maskEdit);
  const layerId = editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0];
  if (!layerId) return undefined;
  const layers = document.artboards.flatMap((artboard) => artboard.layers);
  const selectedLayer = findLayer(layers, layerId);
  if (!selectedLayer || selectedLayer.locked || (!editingMask && !selectedLayer.visible)) return undefined;

  if (editor.tools.activeTool === "feather" && !editingMask && (selectedLayer.type === "image" || selectedLayer.type === "raster")) {
    const maskLayerId = getLayerMask(selectedLayer)?.maskLayerId;
    const attachedMask = maskLayerId ? findLayer(layers, maskLayerId) : undefined;
    const attachedMaskLayer = attachedMask && (attachedMask.type === "image" || attachedMask.type === "raster") ? { ...attachedMask, type: "raster" as const } : undefined;
    const maskAsset = attachedMaskLayer ? document.assets.find((candidate) => candidate.id === attachedMaskLayer.assetId) : undefined;
    if (attachedMaskLayer && maskAsset) return { layer: attachedMaskLayer, asset: maskAsset };

    const sourceAsset = document.assets.find((candidate) => candidate.id === selectedLayer.assetId);
    if (!sourceAsset) return undefined;
    const width = Math.max(1, Math.round(sourceAsset.intrinsicSize.w));
    const height = Math.max(1, Math.round(sourceAsset.intrinsicSize.h));
    const source = opaqueMaskSource(width, height);
    const assetId = crypto.randomUUID() as AssetId;
    const newMaskLayerId = crypto.randomUUID() as LayerId;
    const sourceRect = selectedLayer.sourceRect ?? { x: 0, y: 0, w: width, h: height };
    const maskLayer: RasterLayer = {
      id: newMaskLayerId,
      type: "raster",
      name: `${selectedLayer.name} Mask`,
      visible: true,
      locked: false,
      opacity: 1,
      assetId,
      transform: {
        position: {
          x: selectedLayer.transform.position.x + sourceRect.x * selectedLayer.transform.scale.x,
          y: selectedLayer.transform.position.y + sourceRect.y * selectedLayer.transform.scale.y,
        },
        scale: {
          x: (sourceRect.w * selectedLayer.transform.scale.x) / width,
          y: (sourceRect.h * selectedLayer.transform.scale.y) / height,
        },
        rotation: selectedLayer.transform.rotation,
      },
    };
    const asset = { id: assetId, name: maskLayer.name, mimeType: "image/svg+xml", source, intrinsicSize: { w: width, h: height } };
    return { layer: maskLayer, asset, pendingLayerMask: { targetLayerId: selectedLayer.id, asset, maskLayer } };
  }

  const layer = selectedLayer.type === "raster" ? selectedLayer : undefined;
  if (!layer) return undefined;
  const asset = document.assets.find((candidate) => candidate.id === layer.assetId);
  return asset ? { layer, asset } : undefined;
}

export function updateBrushSession(options: {
  store: AppStore;
  session: BrushSession;
  point: Vec2D;
  color: string;
  size: number;
  hardness: number;
  opacity: number;
  flow: number;
  smoothing: number;
  pressure: number;
  pressureSize: boolean;
  featherRadius?: number;
  featherStrength?: number;
}): BrushSession {
  const state = options.store.getState();
  const asset = state.document.assets.find((candidate) => candidate.id === options.session.assetId);
  if (!asset && !options.session.pendingLayerMask) return options.session;
  const layer = options.session.targetLayer;

  const from = options.session.previousPoint;
  const smoothing = Math.max(0, Math.min(100, options.smoothing)) / 100;
  const follow = 1 - smoothing * 0.85;
  const to = { x: from.x + (options.point.x - from.x) * follow, y: from.y + (options.point.y - from.y) * follow };
  options.session.previousPoint = to;
  options.session.pending = (options.session.pending ?? Promise.resolve())
    .then(async () => {
      if (options.session.cancelled) return;
      if (!(await options.session.ready) || options.session.cancelled) return;

      const assetFrom = documentPointToAssetPoint(from, layer, options.session.width, options.session.height);
      const assetTo = documentPointToAssetPoint(to, layer, options.session.width, options.session.height);
      const size = options.size * (options.pressureSize ? Math.max(0.1, options.pressure) : 1);
      if (options.session.mode === "feather") {
        options.session.featherRadius = options.featherRadius ?? 16;
        drawFeatherSegment(options.session.surface, {
          from: assetFrom,
          to: assetTo,
          size,
          radius: options.featherRadius ?? 16,
          strength: options.featherStrength ?? 65,
        });
      } else {
        drawBrushSegment(options.session.surface, {
          from: assetFrom,
          to: assetTo,
          color: state.editor.maskEdit ? "#ffffff" : options.color,
          size,
          hardness: options.hardness,
          opacity: options.opacity,
          flow: options.flow,
          mode: options.session.mode,
        });
      }

      if (options.session.cancelled) return;
      options.session.changed = true;
      requestBrushStrokePreview({ store: options.store, session: options.session });
    })
    .catch(() => undefined);

  return options.session;
}

export async function commitBrushSession(options: { store: AppStore; session: BrushSession }) {
  await options.session.pending;
  if (options.session.cancelled) return;

  const source = options.session.changed ? brushSurfaceDataUrl(options.session.surface) : undefined;
  if (source) {
    const state = options.store.getState();
    const maskEdit = state.editor.maskEdit;
    const operation = options.session.mode === "feather"
      ? { type: "feather" as const, radius: Math.round(options.session.featherRadius ?? 16) }
      : { type: "paint" as const };
    if (options.session.pendingLayerMask) {
      options.store.dispatch(commandIds.documentAddLayerMask, {
        layerId: options.session.pendingLayerMask.targetLayerId,
        asset: { ...options.session.pendingLayerMask.asset, source, mimeType: "image/png" },
        maskLayer: options.session.pendingLayerMask.maskLayer,
        activeTool: "feather",
      });
    } else if (maskEdit?.kind === "inpaintRegion" && maskEdit.inpaintRegionId && maskEdit.maskAssetId === options.session.assetId) {
      options.store.dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId: maskEdit.inpaintRegionId, source, mimeType: "image/png", operation });
    } else if (maskEdit?.maskLayerId === options.session.layerId) {
      options.store.dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: maskEdit.maskLayerId, source, mimeType: "image/png", operation });
    } else if (isAttachedLayerMask(state.document, options.session.layerId)) {
      options.store.dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: options.session.layerId, source, mimeType: "image/png", operation });
    } else {
      options.store.dispatch(commandIds.documentUpdateAssetSource, { assetId: options.session.assetId, source });
    }
  }
  options.store.dispatch(commandIds.toolSetBrushStrokePreview, undefined);
  closeBrushStrokePreview(options.session);
}

function isPaintTool(tool: ToolState["activeTool"]): tool is "brush" | "eraser" | "feather" {
  return tool === "brush" || tool === "eraser" || tool === "feather";
}

function resolveSelectedFeatherLayer(document: ImageDocument, editor: BrushTargetEditorState) {
  if (isPanInteractionMode(editor.tools.interactionMode)) return undefined;
  const layerId = editor.selection.layerIds[0];
  if (!layerId) return undefined;
  const layer = findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId);
  return layer && (layer.type === "image" || layer.type === "raster") && layer.visible && !layer.locked ? layer : undefined;
}

function isAttachedLayerMask(document: ImageDocument, maskLayerId: string) {
  const stack = document.artboards.flatMap((artboard) => artboard.layers);
  while (stack.length > 0) {
    const layer = stack.pop();
    if (!layer) continue;
    if (getLayerMask(layer)?.maskLayerId === maskLayerId) return true;
    if (layer.type === "group") stack.push(...layer.children);
  }
  return false;
}

function opaqueMaskSource(width: number, height: number) {
  return `data:image/svg+xml,${encodeURIComponent(`<svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${height}" viewBox="0 0 ${width} ${height}"><rect width="${width}" height="${height}" fill="white"/></svg>`)}`;
}

export function cancelBrushSession(options: { store: AppStore; session: BrushSession }) {
  options.session.cancelled = true;
  options.store.dispatch(commandIds.toolSetBrushStrokePreview, undefined);
  closeBrushStrokePreview(options.session);
}

function documentPointToAssetPoint(point: Vec2D, layer: RasterLayer, width: number, height: number): Vec2D {
  const source = layer.sourceRect ?? { x: 0, y: 0, w: width, h: height };
  const destination = {
    x: layer.transform.position.x + source.x * layer.transform.scale.x,
    y: layer.transform.position.y + source.y * layer.transform.scale.y,
    w: source.w * layer.transform.scale.x,
    h: source.h * layer.transform.scale.y,
  };
  const center = { x: destination.x + destination.w / 2, y: destination.y + destination.h / 2 };
  const cos = Math.cos(-layer.transform.rotation);
  const sin = Math.sin(-layer.transform.rotation);
  const dx = point.x - center.x;
  const dy = point.y - center.y;
  const unrotated = { x: center.x + dx * cos - dy * sin, y: center.y + dx * sin + dy * cos };
  return {
    x: source.x + (unrotated.x - destination.x) / Math.max(0.0001, layer.transform.scale.x),
    y: source.y + (unrotated.y - destination.y) / Math.max(0.0001, layer.transform.scale.y),
  };
}

function requestBrushStrokePreview(options: { store: AppStore; session: BrushSession }) {
  if (options.session.cancelled || options.session.previewClosed) return;

  options.session.previewRequested = true;
  if (options.session.previewFrame !== undefined || options.session.previewInFlight) return;

  options.session.previewFrame = scheduleFrame(() => {
    options.session.previewFrame = undefined;
    void publishBrushStrokePreview(options);
  });
}

async function publishBrushStrokePreview(options: { store: AppStore; session: BrushSession }) {
  if (options.session.cancelled || options.session.previewClosed || !options.session.previewRequested) return;

  options.session.previewRequested = false;
  options.session.previewInFlight = true;
  const source = await brushSurfaceObjectUrl(options.session.surface).catch(() => undefined);
  options.session.previewInFlight = false;

  if (!source) {
    if (options.session.previewRequested) requestBrushStrokePreview(options);
    return;
  }

  if (options.session.cancelled || options.session.previewClosed) {
    releaseObjectUrl(source);
    return;
  }

  const previousSource = options.session.previewSource;
  options.session.previewSource = source;
  options.store.dispatch(commandIds.toolSetBrushStrokePreview, {
    layerId: options.session.layerId,
    assetId: options.session.assetId,
    source,
    pendingTargetLayerId: options.session.pendingLayerMask?.targetLayerId,
    intrinsicSize: options.session.pendingLayerMask ? { w: options.session.width, h: options.session.height } : undefined,
  });
  releaseObjectUrl(previousSource);

  if (options.session.previewRequested) requestBrushStrokePreview(options);
}

function closeBrushStrokePreview(session: BrushSession) {
  session.previewClosed = true;
  if (session.previewFrame !== undefined) {
    cancelFrame(session.previewFrame);
    session.previewFrame = undefined;
  }
  if (session.previewSource) {
    releaseObjectUrl(session.previewSource);
    session.previewSource = undefined;
  }
}

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