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

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

import { MAX_HISTORY_STEPS } from './ImageCanvasEditorModel';
import type {
  CanvasGenerationDialogState,
  CanvasHistorySnapshot,
  CanvasLayer,
  CanvasViewport,
  GenerateDialogState,
} from './ImageCanvasEditorTypes';

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

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,
}: {
  refs: CanvasHistoryRefs;
  setters: CanvasHistorySetters;
  resetters: CanvasHistoryResetters;
}) {
  const undoStackRef = useRef<CanvasHistorySnapshot[]>([]);
  const redoStackRef = useRef<CanvasHistorySnapshot[]>([]);
  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) => {
      setters.setLayers(snapshot.layers.map((layer) => ({ ...layer })));
      setters.setViewport({ ...snapshot.viewport });
      setters.setGenerateDialog(
        snapshot.generateDialog
          ? cloneGenerateDialog(snapshot.generateDialog)
          : null,
      );
      setters.setInactiveGenerateDialogs(
        snapshot.inactiveGenerateDialogs.map(cloneCanvasGenerationDialog),
      );
      setters.setSelectedLayerId(snapshot.selectedLayerId);
      setters.setSelectedLayerIds([...snapshot.selectedLayerIds]);
      resetters.setHoveredLayerId(null);
      resetters.setMetadataLayer(null);
      resetters.resetCanvasInteractionState();
    },
    [resetters, setters],
  );

  const captureCanvasHistory = useCallback(
    (options: { clearRedo?: boolean } = {}) => {
      undoStackRef.current = [
        ...undoStackRef.current.slice(-(MAX_HISTORY_STEPS - 1)),
        getCanvasHistorySnapshot(),
      ];
      if (options.clearRedo !== false) {
        redoStackRef.current = [];
      }
      setHistoryVersion((version) => version + 1);
    },
    [getCanvasHistorySnapshot],
  );

  const undoCanvasChange = useCallback(() => {
    const previousSnapshot = undoStackRef.current.at(-1);
    if (!previousSnapshot) {
      return;
    }
    undoStackRef.current = undoStackRef.current.slice(0, -1);
    redoStackRef.current = [
      ...redoStackRef.current.slice(-(MAX_HISTORY_STEPS - 1)),
      getCanvasHistorySnapshot(),
    ];
    restoreCanvasHistorySnapshot(previousSnapshot);
    setHistoryVersion((version) => version + 1);
  }, [getCanvasHistorySnapshot, restoreCanvasHistorySnapshot]);

  const redoCanvasChange = useCallback(() => {
    const nextSnapshot = redoStackRef.current.at(-1);
    if (!nextSnapshot) {
      return;
    }
    redoStackRef.current = redoStackRef.current.slice(0, -1);
    undoStackRef.current = [
      ...undoStackRef.current.slice(-(MAX_HISTORY_STEPS - 1)),
      getCanvasHistorySnapshot(),
    ];
    restoreCanvasHistorySnapshot(nextSnapshot);
    setHistoryVersion((version) => version + 1);
  }, [getCanvasHistorySnapshot, restoreCanvasHistorySnapshot]);

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