useSuperActionState

import { type SubmissionResult, useForm } from "@conform-to/react";
import { parseWithZod } from "@conform-to/zod";
import { useActionState, useCallback } from "react";
import type { TypeOf, ZodTypeAny } from "zod";

type State = SubmissionResult<string[]> | null | undefined;

interface Callbacks<T, R = unknown> {
  onStart?: () => R;
  onEnd?: (reference: R) => void;
  onSuccess?: (result: T) => void;
  onError?: (result: T) => void;
}

const useActionStateWithCallbacks = <Reference,>(
  action: (state: Awaited<State>, payload: FormData) => Promise<State>,
  initialState: Awaited<State>,
  callbacks: Callbacks<Awaited<State>, Reference>,
) => {
  const _action = useCallback(
    async (...args: Parameters<typeof action>) => {
      const promise = action(...args);

      const reference = callbacks.onStart?.();

      const result = await promise;

      if (reference) {
        callbacks.onEnd?.(reference);
      }

      if (result?.status === "success") {
        callbacks.onSuccess?.(result);
      }

      if (result?.status === "error") {
        callbacks.onError?.(result);
      }

      return promise;
    },
    [action, callbacks],
  );

  return useActionState(_action, initialState);
};

export const useSuperActionState = <
  Schema extends ZodTypeAny,
  Reference = unknown,
>(
  action: (state: Awaited<State>, payload: FormData) => Promise<State>,
  options: {
    schema?: Schema;
    additionalData?: Record<string, string>;
    callbacks?: Callbacks<Awaited<State>, Reference>;
  } & Parameters<typeof useForm<TypeOf<Schema>>>[0] = {},
): [
  state: ReturnType<typeof useForm<TypeOf<Schema>>>,
  dispatch: (formData: FormData) => void,
  isPending: boolean,
] => {
  const { schema, additionalData = {}, callbacks, ...useFormOptions } = options;

  const _action = (state: State, payload: FormData) => {
    for (const [name, value] of Object.entries(additionalData)) {
      payload.append(name, value);
    }
    return action(state, payload);
  };

  const [lastResult, dispatch, isPending] = callbacks
    ? useActionStateWithCallbacks(_action, null, callbacks)
    : useActionState(_action, null);

  const [form, fields] = useForm({
    lastResult,
    onValidate: schema
      ? ({ formData }) => {
          return parseWithZod(formData, { schema });
        }
      : undefined,
    ...useFormOptions,
  });

  return [[form, fields], dispatch, isPending];
};

useEffect(() => {
    const preventDefault = (event: Event) => {
      // Make sure the reset event is dispatched on the corresponding form element
      if (event.target === document.forms.namedItem(form.id)) {
        // Tell Conform to ignore the form reset event
        event.preventDefault();
      }
    };

    document.addEventListener("reset", preventDefault, true);

    return () => {
      document.removeEventListener("reset", preventDefault, true);
    };
}, [form.id]);