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

import {
  type Dispatch,
  type SetStateAction,
  useCallback,
  useMemo,
  useRef,
  useState,
} from 'react';

import { isCanvasGenerationDialog } from './ImageCanvasGenerationModel';
import type {
  CanvasGenerationDialogState,
  GenerateDialogState,
} from './ImageCanvasEditorTypes';

type CanvasGenerationDialogUpdater = (
  dialog: CanvasGenerationDialogState,
) => CanvasGenerationDialogState | null;

function withGenerationTimestamps<T extends GenerateDialogState | null>(
  nextDialog: T,
  previousDialog?: GenerateDialogState | null,
): T {
  if (!nextDialog) {
    return nextDialog;
  }
  if (nextDialog.status !== 'generating') {
    const shouldMarkFinished =
      nextDialog.status === 'failed' || Boolean(nextDialog.generatedLayerId);
    return {
      ...nextDialog,
      generationFinishedAt: shouldMarkFinished
        ? (nextDialog.generationFinishedAt ??
          previousDialog?.generationFinishedAt ??
          Date.now())
        : nextDialog.generationFinishedAt,
    };
  }
  return {
    ...nextDialog,
    generationStartedAt:
      nextDialog.generationStartedAt ??
      previousDialog?.generationStartedAt ??
      Date.now(),
    generationFinishedAt: undefined,
  };
}

export function useCanvasGenerationDialogs({
  onActivate,
}: {
  onActivate?: () => void;
} = {}) {
  const generationDialogCounterRef = useRef(0);
  const generateDialogRef = useRef<GenerateDialogState | null>(null);
  const inactiveGenerateDialogsRef = useRef<CanvasGenerationDialogState[]>([]);
  const [generateDialog, setGenerateDialogState] =
    useState<GenerateDialogState | null>(null);
  const [inactiveGenerateDialogs, setInactiveGenerateDialogs] = useState<
    CanvasGenerationDialogState[]
  >([]);

  generateDialogRef.current = generateDialog;
  inactiveGenerateDialogsRef.current = inactiveGenerateDialogs;

  const activeCanvasGenerationDialog = isCanvasGenerationDialog(generateDialog)
    ? generateDialog
    : null;

  const setGenerateDialog = useCallback<
    Dispatch<SetStateAction<GenerateDialogState | null>>
  >((nextDialogOrUpdater) => {
    const currentDialog = generateDialogRef.current;
    const nextDialog =
      typeof nextDialogOrUpdater === 'function'
        ? nextDialogOrUpdater(currentDialog)
        : nextDialogOrUpdater;
    const nextDialogWithTimestamps = withGenerationTimestamps(
      nextDialog,
      currentDialog,
    );
    generateDialogRef.current = nextDialogWithTimestamps;
    setGenerateDialogState(nextDialogWithTimestamps);
  }, []);
  const canvasGenerationDialogs = useMemo(
    () =>
      activeCanvasGenerationDialog
        ? [...inactiveGenerateDialogs, activeCanvasGenerationDialog]
        : inactiveGenerateDialogs,
    [activeCanvasGenerationDialog, inactiveGenerateDialogs],
  );

  const createGenerationDialogId = useCallback(() => {
    generationDialogCounterRef.current += 1;
    return `generation-dialog-${generationDialogCounterRef.current}`;
  }, []);

  const archiveActiveCanvasGenerationDialog = useCallback(() => {
    const currentDialog = generateDialogRef.current;
    if (!isCanvasGenerationDialog(currentDialog)) {
      return;
    }
    const nextInactiveDialogs = inactiveGenerateDialogsRef.current.some(
      (dialog) => dialog.id === currentDialog.id,
    )
      ? inactiveGenerateDialogsRef.current
      : [
          ...inactiveGenerateDialogsRef.current,
          {
            ...currentDialog,
            composerOpen: false,
          },
        ];
    inactiveGenerateDialogsRef.current = nextInactiveDialogs;
    setInactiveGenerateDialogs(nextInactiveDialogs);
  }, []);

  const openCanvasGenerationDialog = useCallback(
    (dialog: Omit<CanvasGenerationDialogState, 'id'>) => {
      const currentDialog = generateDialogRef.current;
      if (isCanvasGenerationDialog(currentDialog)) {
        inactiveGenerateDialogsRef.current =
          inactiveGenerateDialogsRef.current.some(
            (existingDialog) => existingDialog.id === currentDialog.id,
          )
            ? inactiveGenerateDialogsRef.current
            : [
                ...inactiveGenerateDialogsRef.current,
                {
                  ...currentDialog,
                  composerOpen: false,
                },
              ];
      }
      archiveActiveCanvasGenerationDialog();
      const id = createGenerationDialogId();
      const nextDialog = withGenerationTimestamps({
        ...dialog,
        id,
      });
      generateDialogRef.current = nextDialog;
      setGenerateDialogState(nextDialog);
      return id;
    },
    [archiveActiveCanvasGenerationDialog, createGenerationDialogId],
  );

  const updateCanvasGenerationDialogById = useCallback(
    (dialogId: string, updater: CanvasGenerationDialogUpdater) => {
      const currentDialogRef = generateDialogRef.current;
      if (
        isCanvasGenerationDialog(currentDialogRef) &&
        currentDialogRef.id === dialogId
      ) {
        generateDialogRef.current = withGenerationTimestamps(
          updater(currentDialogRef),
          currentDialogRef,
        );
      }
      inactiveGenerateDialogsRef.current =
        inactiveGenerateDialogsRef.current.flatMap((dialog) => {
          if (dialog.id !== dialogId) {
            return [dialog];
          }
          const nextDialog = withGenerationTimestamps(updater(dialog), dialog);
          return nextDialog ? [nextDialog] : [];
        });
      setGenerateDialogState((currentDialog) => {
        if (
          !isCanvasGenerationDialog(currentDialog) ||
          currentDialog.id !== dialogId
        ) {
          return currentDialog;
        }
        return withGenerationTimestamps(updater(currentDialog), currentDialog);
      });
      setInactiveGenerateDialogs((currentDialogs) =>
        currentDialogs.flatMap((dialog) => {
          if (dialog.id !== dialogId) {
            return [dialog];
          }
          const nextDialog = withGenerationTimestamps(updater(dialog), dialog);
          return nextDialog ? [nextDialog] : [];
        }),
      );
    },
    [],
  );

  const removeCanvasGenerationDialogById = useCallback(
    (dialogId: string) => {
      updateCanvasGenerationDialogById(dialogId, () => null);
    },
    [updateCanvasGenerationDialogById],
  );

  const hasCanvasGenerationDialogById = useCallback((dialogId: string) => {
    const currentDialog = generateDialogRef.current;
    return [
      ...(isCanvasGenerationDialog(currentDialog) ? [currentDialog] : []),
      ...inactiveGenerateDialogsRef.current,
    ].some((dialog) => dialog.id === dialogId);
  }, []);

  const activateCanvasGenerationDialog = useCallback(
    (targetDialog: CanvasGenerationDialogState) => {
      const currentDialog = generateDialogRef.current;
      const nextInactiveDialogs = inactiveGenerateDialogsRef.current.filter(
        (dialog) => dialog.id !== targetDialog.id,
      );
      if (
        isCanvasGenerationDialog(currentDialog) &&
        currentDialog.id !== targetDialog.id
      ) {
        nextInactiveDialogs.push({
          ...currentDialog,
          composerOpen: false,
        });
      }
      const nextActiveDialog = {
        ...targetDialog,
        composerOpen: true,
      };
      inactiveGenerateDialogsRef.current = nextInactiveDialogs;
      generateDialogRef.current = nextActiveDialog;
      setInactiveGenerateDialogs(nextInactiveDialogs);
      setGenerateDialogState(nextActiveDialog);
      onActivate?.();
    },
    [onActivate],
  );

  const restoreCanvasGenerationDialogs = useCallback(
    (dialogs: CanvasGenerationDialogState[]) => {
      const nextCounter = dialogs.reduce((maxCounter, dialog) => {
        const match = /^generation-dialog-(\d+)$/u.exec(dialog.id);
        const numericId = match ? Number.parseInt(match[1] ?? '0', 10) : 0;
        return Math.max(maxCounter, Number.isFinite(numericId) ? numericId : 0);
      }, 0);
      generationDialogCounterRef.current = Math.max(
        generationDialogCounterRef.current,
        nextCounter,
      );
      const activeDialog =
        [...dialogs].reverse().find((dialog) => dialog.composerOpen !== false) ??
        dialogs[dialogs.length - 1] ??
        null;
      const nextActiveDialog = activeDialog
        ? {
            ...activeDialog,
            composerOpen: activeDialog.composerOpen !== false,
          }
        : null;
      const nextInactiveDialogs = dialogs
        .filter((dialog) => dialog.id !== activeDialog?.id)
        .map((dialog) => ({
          ...dialog,
          composerOpen: false,
        }));
      generateDialogRef.current = nextActiveDialog;
      inactiveGenerateDialogsRef.current = nextInactiveDialogs;
      setGenerateDialogState(nextActiveDialog);
      setInactiveGenerateDialogs(nextInactiveDialogs);
    },
    [],
  );

  const removeCanvasGenerationDialogsByLayerId = useCallback(
    (targetLayerId: string) => {
      const keepDialog = (dialog: CanvasGenerationDialogState) =>
        dialog.sourceLayerId !== targetLayerId &&
        dialog.generatedLayerId !== targetLayerId;
      const currentDialog = generateDialogRef.current;
      const nextActiveDialog =
        isCanvasGenerationDialog(currentDialog) && !keepDialog(currentDialog)
          ? null
          : currentDialog;
      const nextInactiveDialogs =
        inactiveGenerateDialogsRef.current.filter(keepDialog);
      generateDialogRef.current = nextActiveDialog;
      inactiveGenerateDialogsRef.current = nextInactiveDialogs;
      setGenerateDialogState(nextActiveDialog);
      setInactiveGenerateDialogs(nextInactiveDialogs);
    },
    [],
  );

  const getGeneratingDialogPlaceholder = useCallback(
    (dialog: GenerateDialogState) => {
      const currentDialog = generateDialogRef.current;
      if (dialog.id) {
        const latestDialog = [
          ...(isCanvasGenerationDialog(currentDialog) ? [currentDialog] : []),
          ...inactiveGenerateDialogsRef.current,
        ].find((candidateDialog) => candidateDialog.id === dialog.id);
        if (latestDialog?.status === 'generating') {
          return latestDialog.placeholder ?? dialog.placeholder;
        }
      }
      if (
        currentDialog?.mode === dialog.mode &&
        (!dialog.id || currentDialog.id === dialog.id) &&
        currentDialog.status === 'generating'
      ) {
        return currentDialog.placeholder ?? dialog.placeholder;
      }
      return dialog.placeholder;
    },
    [],
  );

  return {
    generateDialog,
    setGenerateDialog,
    generateDialogRef,
    inactiveGenerateDialogs,
    setInactiveGenerateDialogs,
    inactiveGenerateDialogsRef,
    activeCanvasGenerationDialog,
    canvasGenerationDialogs,
    archiveActiveCanvasGenerationDialog,
    openCanvasGenerationDialog,
    updateCanvasGenerationDialogById,
    removeCanvasGenerationDialogById,
    hasCanvasGenerationDialogById,
    activateCanvasGenerationDialog,
    restoreCanvasGenerationDialogs,
    removeCanvasGenerationDialogsByLayerId,
    getGeneratingDialogPlaceholder,
  };
}