Genarrative/src/components/image-editor/useCanvasHistory.ts

import { type RefObject, useCallback, useRef, useState } from 'react';

import { MAX_HISTORY_STEPS } from './ImageCanvasEditorModel';
import type {
  CanvasGenerationDialogState,
  CanvasHistoryAction,
  CanvasHistoryApplyResult,
  CanvasHistoryEntry,
  CanvasHistorySnapshot,
  CanvasLayer,
  CanvasViewport,
  GenerateDialogState,
} from './ImageCanvasEditorTypes';
import {
  canRestoreCanvasHistorySnapshotWithoutContentLoss,
  isProtectedCanvasHistoryAction,
  mergeCanvasHistorySnapshotForRestore,
} from './ImageCanvasHistoryModel';

type CanvasHistoryRefs = {
  layersRef: RefObject<CanvasLayer[]>;
  viewportRef: RefObject<CanvasViewport>;
  generateDialogRef: RefObject<GenerateDialogState | null>;
  inactiveGenerateDialogsRef: RefObject<CanvasGenerationDialogState[]>;
  selectedLayerIdRef: RefObject<string | null>;
  selectedLayerIdsRef: RefObject<string[]>;
};

type CanvasHistorySetters = {
  setLayers: (layers: CanvasLayer[]) => void;
  setViewport: (viewport: CanvasViewport) => void;
  setGenerateDialog: (dialog: GenerateDialogState | null) => void;
  setInactiveGenerateDialogs: (dialogs: CanvasGenerationDialogState[]) => void;
  setSelectedLayerId: (layerId: string | null) => void;
  setSelectedLayerIds: (layerIds: string[]) => void;
};

type CanvasHistoryResetters = {
  setHoveredLayerId: (layerId: string | null) => void;
  setMetadataLayer: (layer: CanvasLayer | null) => void;
  resetCanvasInteractionState: () => void;
};

type CanvasHistoryEffects = {
  persistRestoredAssetKind?: (layer: CanvasLayer) => void;
};

function cloneGenerateDialog(dialog: GenerateDialogState): GenerateDialogState {
  return {
    ...dialog,
    placeholder: dialog.placeholder ? { ...dialog.placeholder } : undefined,
  };
}

function cloneCanvasGenerationDialog(
  dialog: CanvasGenerationDialogState,
): CanvasGenerationDialogState {
  return {
    ...dialog,
    placeholder: dialog.placeholder ? { ...dialog.placeholder } : undefined,
  };
}

export function useCanvasHistory({
  refs,
  setters,
  resetters,
  effects,
}: {
  refs: CanvasHistoryRefs;
  setters: CanvasHistorySetters;
  resetters: CanvasHistoryResetters;
  effects?: CanvasHistoryEffects;
}) {
  const undoStackRef = useRef<CanvasHistoryEntry[]>([]);
  const redoStackRef = useRef<CanvasHistoryEntry[]>([]);
  const [historyVersion, setHistoryVersion] = useState(0);

  const getCanvasHistorySnapshot = useCallback(
    (): CanvasHistorySnapshot => ({
      layers: refs.layersRef.current.map((layer) => ({ ...layer })),
      viewport: { ...refs.viewportRef.current },
      generateDialog: refs.generateDialogRef.current
        ? cloneGenerateDialog(refs.generateDialogRef.current)
        : null,
      inactiveGenerateDialogs: refs.inactiveGenerateDialogsRef.current.map(
        cloneCanvasGenerationDialog,
      ),
      selectedLayerId: refs.selectedLayerIdRef.current,
      selectedLayerIds: [...refs.selectedLayerIdsRef.current],
    }),
    [refs],
  );

  const restoreCanvasHistorySnapshot = useCallback(
    (snapshot: CanvasHistorySnapshot, action?: CanvasHistoryAction) => {
      const currentSnapshot = getCanvasHistorySnapshot();
      const mergedSnapshot = mergeCanvasHistorySnapshotForRestore({
        current: currentSnapshot,
        target: snapshot,
      });
      setters.setLayers(mergedSnapshot.layers.map((layer) => ({ ...layer })));
      setters.setViewport({ ...mergedSnapshot.viewport });
      setters.setGenerateDialog(
        mergedSnapshot.generateDialog
          ? cloneGenerateDialog(mergedSnapshot.generateDialog)
          : null,
      );
      setters.setInactiveGenerateDialogs(
        mergedSnapshot.inactiveGenerateDialogs.map(cloneCanvasGenerationDialog),
      );
      setters.setSelectedLayerId(mergedSnapshot.selectedLayerId);
      setters.setSelectedLayerIds([...mergedSnapshot.selectedLayerIds]);
      resetters.setHoveredLayerId(null);
      resetters.setMetadataLayer(null);
      resetters.resetCanvasInteractionState();

      if (action?.type === 'change-asset-kind') {
        const currentLayerById = new Map(
          currentSnapshot.layers.map((layer) => [layer.id, layer] as const),
        );
        for (const layer of mergedSnapshot.layers) {
          const currentLayer = currentLayerById.get(layer.id);
          if (currentLayer && currentLayer.assetKind !== layer.assetKind) {
            effects?.persistRestoredAssetKind?.({ ...layer });
          }
        }
      }
    },
    [effects, getCanvasHistorySnapshot, resetters, setters],
  );

  const captureCanvasHistory = useCallback(
    (
      action: CanvasHistoryAction,
      options: { snapshot?: CanvasHistorySnapshot } = {},
    ) => {
      undoStackRef.current = [
        ...undoStackRef.current.slice(-(MAX_HISTORY_STEPS - 1)),
        {
          snapshot: options.snapshot ?? getCanvasHistorySnapshot(),
          action,
          createdAt: Date.now(),
        },
      ];
      redoStackRef.current = [];
      setHistoryVersion((version) => version + 1);
    },
    [getCanvasHistorySnapshot],
  );

  const discardCanvasHistoryEntriesContainingLayer = useCallback(
    (matchesLayer: (layer: CanvasLayer) => boolean) => {
      const keepEntry = (entry: CanvasHistoryEntry) =>
        !entry.snapshot.layers.some(matchesLayer);
      const nextUndoStack = undoStackRef.current.filter(keepEntry);
      const nextRedoStack = redoStackRef.current.filter(keepEntry);
      if (
        nextUndoStack.length === undoStackRef.current.length &&
        nextRedoStack.length === redoStackRef.current.length
      ) {
        return;
      }
      undoStackRef.current = nextUndoStack;
      redoStackRef.current = nextRedoStack;
      setHistoryVersion((version) => version + 1);
    },
    [],
  );

  const undoCanvasChange = useCallback((): CanvasHistoryApplyResult => {
    const previousEntry = undoStackRef.current.at(-1);
    if (!previousEntry) {
      return { status: 'empty' };
    }
    const currentSnapshot = getCanvasHistorySnapshot();
    if (
      isProtectedCanvasHistoryAction(previousEntry.action) ||
      !canRestoreCanvasHistorySnapshotWithoutContentLoss({
        current: currentSnapshot,
        target: previousEntry.snapshot,
      })
    ) {
      return {
        status: 'blocked',
        action: previousEntry.action,
        reason: 'content-may-disappear',
      };
    }
    undoStackRef.current = undoStackRef.current.slice(0, -1);
    redoStackRef.current = [
      ...redoStackRef.current.slice(-(MAX_HISTORY_STEPS - 1)),
      {
        snapshot: currentSnapshot,
        action: previousEntry.action,
        createdAt: Date.now(),
      },
    ];
    restoreCanvasHistorySnapshot(previousEntry.snapshot, previousEntry.action);
    setHistoryVersion((version) => version + 1);
    return { status: 'success', action: previousEntry.action };
  }, [getCanvasHistorySnapshot, restoreCanvasHistorySnapshot]);

  const redoCanvasChange = useCallback((): CanvasHistoryApplyResult => {
    const nextEntry = redoStackRef.current.at(-1);
    if (!nextEntry) {
      return { status: 'empty' };
    }
    const currentSnapshot = getCanvasHistorySnapshot();
    if (
      !canRestoreCanvasHistorySnapshotWithoutContentLoss({
        current: currentSnapshot,
        target: nextEntry.snapshot,
      })
    ) {
      return {
        status: 'blocked',
        action: nextEntry.action,
        reason: 'content-may-disappear',
      };
    }
    redoStackRef.current = redoStackRef.current.slice(0, -1);
    undoStackRef.current = [
      ...undoStackRef.current.slice(-(MAX_HISTORY_STEPS - 1)),
      {
        snapshot: currentSnapshot,
        action: nextEntry.action,
        createdAt: Date.now(),
      },
    ];
    restoreCanvasHistorySnapshot(nextEntry.snapshot, nextEntry.action);
    setHistoryVersion((version) => version + 1);
    return { status: 'success', action: nextEntry.action };
  }, [getCanvasHistorySnapshot, restoreCanvasHistorySnapshot]);

  return {
    canUndo: undoStackRef.current.length > 0,
    canRedo: redoStackRef.current.length > 0,
    historyVersion,
    getCanvasHistorySnapshot,
    restoreCanvasHistorySnapshot,
    captureCanvasHistory,
    discardCanvasHistoryEntriesContainingLayer,
    undoCanvasChange,
    redoCanvasChange,
  };
}