weseek__growi/apps/app/src/features/page-tree/client/states/page-tree-update.ts

import { useCallback, useEffect } from 'react';
import type { TreeInstance } from '@headless-tree/core';
import { atom, useAtomValue, useSetAtom } from 'jotai';

import { ROOT_PAGE_VIRTUAL_ID } from '../../constants';
import { invalidatePageTreeChildren } from '../services';

// Update generation number
const generationAtom = atom<number>(1);

// Array of IDs for last updated items
// null is a special value meaning full tree update
const lastUpdatedItemIdsAtom = atom<string[] | null>(null);

// Read-only hooks
export const usePageTreeInformationGeneration = () =>
  useAtomValue(generationAtom);

export const usePageTreeInformationLastUpdatedItemIds = () =>
  useAtomValue(lastUpdatedItemIdsAtom);

// Hook for notifying tree updates
export const usePageTreeInformationUpdate = () => {
  const setGeneration = useSetAtom(generationAtom);
  const setLastUpdatedIds = useSetAtom(lastUpdatedItemIdsAtom);

  // Notify update for specific items
  const notifyUpdateItems = useCallback(
    (itemIds: string[]) => {
      setLastUpdatedIds(itemIds);
      setGeneration((prev) => prev + 1);
    },
    [setGeneration, setLastUpdatedIds],
  );

  // Notify update for all trees
  const notifyUpdateAllTrees = useCallback(() => {
    setLastUpdatedIds(null);
    setGeneration((prev) => prev + 1);
  }, [setGeneration, setLastUpdatedIds]);

  return {
    notifyUpdateItems,
    notifyUpdateAllTrees,
  };
};

export const usePageTreeRevalidationEffect = (
  tree: TreeInstance<unknown>,
  generation: number,
  opts?: { onRevalidated?: () => void },
) => {
  const globalGeneration = useAtomValue(generationAtom);
  const globalLastUpdatedItemIds = useAtomValue(lastUpdatedItemIdsAtom);

  const { getItemInstance, rebuildTree } = tree;

  useEffect(() => {
    if (globalGeneration <= generation) return; // Already up to date

    // Determine update scope
    const shouldUpdateAll = globalLastUpdatedItemIds == null;

    if (shouldUpdateAll) {
      // Full tree update: clear all cache and refetch from root
      invalidatePageTreeChildren();
      const root = getItemInstance(ROOT_PAGE_VIRTUAL_ID);
      root?.invalidateChildrenIds(true);
    } else {
      // Partial update: clear cache for specified items and refetch children
      invalidatePageTreeChildren(globalLastUpdatedItemIds);
      globalLastUpdatedItemIds.forEach((itemId) => {
        const item = getItemInstance(itemId);
        // Invalidate children to refresh child list
        item?.invalidateChildrenIds(true);
      });
    }

    // Rebuild tree after a short delay to allow async data fetching to complete
    // This ensures isItemFolder is re-evaluated with fresh children data
    const timeoutId = setTimeout(() => {
      rebuildTree();
    }, 100);

    opts?.onRevalidated?.();

    return () => clearTimeout(timeoutId);
  }, [
    globalGeneration,
    generation,
    getItemInstance,
    globalLastUpdatedItemIds,
    rebuildTree,
    opts,
  ]);
};