image-studio/view/canvas/brush.ts

import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
import type { Layer } from "@core/layer";
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";

export type BrushSession = {
  layerId: string;
  assetId: string;
  width: number;
  height: number;
  canvas: HTMLCanvasElement;
  context: CanvasRenderingContext2D;
  ready: Promise<boolean>;
  previousPoint: Vec2D;
  mode: "brush" | "eraser";
  changed?: boolean;
  pending?: Promise<void>;
  cancelled?: boolean;
  previewClosed?: boolean;
  previewRequested?: boolean;
  previewInFlight?: boolean;
  previewFrame?: number;
  previewSource?: string;
};

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

export function beginBrushSession(document: ImageDocument, editor: BrushTargetEditorState, point: Vec2D): BrushSession | undefined {
  const layer = resolveBrushTargetLayer(document, editor);
  if (!layer || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined;

  const asset = document.assets.find((candidate) => candidate.id === layer.assetId);
  if (!asset) return undefined;

  const canvas = globalThis.document.createElement("canvas");
  canvas.width = Math.max(1, Math.round(asset.intrinsicSize.w));
  canvas.height = Math.max(1, Math.round(asset.intrinsicSize.h));
  const context = canvas.getContext("2d");
  if (!context) return undefined;

  const session: BrushSession = {
    layerId: layer.id,
    assetId: layer.assetId,
    width: canvas.width,
    height: canvas.height,
    canvas,
    context,
    ready: Promise.resolve(false),
    previousPoint: point,
    mode: editor.tools.activeTool,
  };
  session.ready = initializeBrushSession(session, asset.source).catch(() => false);
  return session;
}

export function canPreviewBrush(document: ImageDocument, editor: BrushTargetEditorState): boolean {
  return Boolean(resolveBrushTargetLayer(document, editor));
}

export function brushUnavailableHint(document: ImageDocument, editor: BrushTargetEditorState): string | undefined {
  if (isPanInteractionMode(editor.tools.interactionMode) || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined;
  if (resolveBrushTargetLayer(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") 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 resolveBrushTargetLayer(document: ImageDocument, editor: BrushTargetEditorState): RasterLayer | undefined {
  if (isPanInteractionMode(editor.tools.interactionMode) || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined;
  const editingMask = Boolean(editor.maskEdit);
  const layerId = editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0];
  if (!layerId) return undefined;
  const layer = findRasterLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId);
  if (!layer || layer.locked || (!editingMask && !layer.visible)) return undefined;
  return layer;
}

export function updateBrushSession(options: {
  store: AppStore;
  session: BrushSession;
  point: Vec2D;
  color: string;
  size: number;
  hardness: number;
}): BrushSession {
  const state = options.store.getState();
  const layer = findRasterLayer(state.document.artboards.flatMap((artboard) => artboard.layers), options.session.layerId);
  if (!layer || layer.assetId !== options.session.assetId) return options.session;

  const from = options.session.previousPoint;
  const to = options.point;
  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;

      drawStrokeSegment({
        context: options.session.context,
        from: documentPointToAssetPoint(from, layer, options.session.width, options.session.height),
        to: documentPointToAssetPoint(to, layer, options.session.width, options.session.height),
        color: state.editor.maskEdit ? "#ffffff" : options.color,
        size: options.size,
        hardness: options.hardness,
        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 ? canvasToDataUrl(options.session.canvas) : undefined;
  if (source) {
    const state = options.store.getState();
    const maskEdit = state.editor.maskEdit;
    if (maskEdit?.maskLayerId === options.session.layerId) {
      options.store.dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: maskEdit.maskLayerId, source, mimeType: "image/png", operation: { type: "paint" } });
    } else {
      options.store.dispatch(commandIds.documentUpdateAssetSource, { assetId: options.session.assetId, source });
    }
  }
  options.store.dispatch(commandIds.toolSetBrushStrokePreview, undefined);
  closeBrushStrokePreview(options.session);
}

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 {
  return {
    x: ((point.x - layer.transform.position.x) / Math.max(0.0001, layer.transform.scale.x) / width) * width,
    y: ((point.y - layer.transform.position.y) / Math.max(0.0001, layer.transform.scale.y) / height) * height,
  };
}

async function initializeBrushSession(session: BrushSession, source: string) {
  const image = await loadImage(source);
  if (session.cancelled) return false;

  session.context.clearRect(0, 0, session.width, session.height);
  session.context.drawImage(image, 0, 0, session.width, session.height);
  return true;
}

function drawStrokeSegment(options: {
  context: CanvasRenderingContext2D;
  from: Vec2D;
  to: Vec2D;
  color: string;
  size: number;
  hardness: number;
  mode: "brush" | "eraser";
}) {
  const { context } = options;
  const hardness = Math.max(0, Math.min(100, options.hardness)) / 100;
  context.save();
  context.globalCompositeOperation = options.mode === "eraser" ? "destination-out" : "source-over";
  context.strokeStyle = options.color;
  context.shadowColor = options.mode === "eraser" ? "rgba(0,0,0,1)" : options.color;
  context.shadowBlur = (1 - hardness) * options.size;
  context.lineWidth = options.size;
  context.lineCap = "round";
  context.lineJoin = "round";
  context.beginPath();
  context.moveTo(options.from.x, options.from.y);
  context.lineTo(options.to.x, options.to.y);
  context.stroke();
  context.restore();
}

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 = requestAnimationFrame(() => {
    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 canvasToObjectUrl(options.session.canvas).catch(() => undefined);
  options.session.previewInFlight = false;

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

  if (options.session.cancelled || options.session.previewClosed) {
    URL.revokeObjectURL(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 });
  if (previousSource) URL.revokeObjectURL(previousSource);

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

function closeBrushStrokePreview(session: BrushSession) {
  session.previewClosed = true;
  if (session.previewFrame !== undefined) {
    cancelAnimationFrame(session.previewFrame);
    session.previewFrame = undefined;
  }
  if (session.previewSource) {
    URL.revokeObjectURL(session.previewSource);
    session.previewSource = undefined;
  }
}

function canvasToObjectUrl(canvas: HTMLCanvasElement) {
  return new Promise<string | undefined>((resolve) => {
    canvas.toBlob((blob) => {
      resolve(blob ? URL.createObjectURL(blob) : undefined);
    }, "image/png");
  });
}

function canvasToDataUrl(canvas: HTMLCanvasElement) {
  try {
    return canvas.toDataURL("image/png");
  } catch {
    return undefined;
  }
}

function loadImage(source: string) {
  return new Promise<HTMLImageElement>((resolve, reject) => {
    const image = new Image();
    image.onload = () => resolve(image);
    image.onerror = () => reject(new Error("Failed to load raster layer"));
    image.src = source;
  });
}

function findRasterLayer(layers: Layer[], layerId: string): RasterLayer | undefined {
  const layer = findLayer(layers, layerId);
  return layer?.type === "raster" ? layer : 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;
}