import { useMutation, useQueryClient } from '@tanstack/react-query';
import type { MutationFunction, QueryKey } from '@tanstack/react-query';

type InvalidatingMutationOptions<TData, TVariables> = {
  readonly mutationFn: MutationFunction<TData, TVariables>;
  readonly invalidateQueryKey?: QueryKey;
  readonly invalidateQueryKeys?: readonly QueryKey[];
  readonly onSuccess?: (data: TData, variables: TVariables) => void | Promise<void>;
};

export function useInvalidatingMutation<TData = unknown, TVariables = void>({
  mutationFn,
  invalidateQueryKey,
  invalidateQueryKeys = [],
  onSuccess,
}: InvalidatingMutationOptions<TData, TVariables>) {
  const queryClient = useQueryClient();
  const queryKeys = invalidateQueryKey
    ? [invalidateQueryKey, ...invalidateQueryKeys]
    : invalidateQueryKeys;

  return useMutation({
    mutationFn,
    onSuccess: async (data, variables) => {
      await Promise.all(queryKeys.map((queryKey) => (
        queryClient.invalidateQueries({ queryKey })
      )));
      await onSuccess?.(data, variables);
    },
  });
}
