import { useQueryClient } from "@tanstack/react-query";
import { useServerMutation } from "./use-server-mutation";

export function useOptimisticMutation<TData, TVariables = unknown>(
  options: Parameters<typeof useServerMutation<TData, TVariables>>[0] & {
    queryKey: readonly (string | number)[];
    optimisticUpdate?: (oldData: any, variables: TVariables) => any;
  }
) {
  const queryClient = useQueryClient();
  const { queryKey, optimisticUpdate, ...mutationOptions } = options;

  return useServerMutation({
    ...mutationOptions,
    onMutate: async (variables) => {
      await queryClient.cancelQueries({ queryKey });

      const previousData = queryClient.getQueryData(queryKey);

      if (optimisticUpdate && previousData) {
        queryClient.setQueryData(queryKey, (old: any) => optimisticUpdate(old, variables));
      }

      return { previousData };
    },
    onError: (error, variables, onMutateResult, context) => {
      if (context && typeof context === 'object' && 'previousData' in context && context.previousData) {
        queryClient.setQueryData(queryKey, context.previousData);
      }
      mutationOptions.onError?.(error, variables, onMutateResult, context);
    },
    onSettled: (data, error, variables, onMutateResult, context) => {
      queryClient.invalidateQueries({ queryKey });
      mutationOptions.onSettled?.(data, error, variables, onMutateResult, context);
    }
  });
}