useMediatedState

/**
 * @template T
 * @callback MediatorUnary
 * @param {T} newState
 * @returns {T}
 */

/**
 * @template T
 * @callback MediatorDispatch
 * @param {T} newState
 * @param {import('react').Dispatch<import('react').SetStateAction<T>>} dispatch
 * @returns {void}
 */

/**
 * @template T
 * @param {MediatorUnary<T> | MediatorDispatch<T>} mediator
 * @param {T | (() => T) | undefined} [initialState]
 * @returns {[T | undefined, import('react').Dispatch<import('react').SetStateAction<T | undefined>>]}
 */
/**
 * mediator: A function that receives the new state and returns the transformed state. 
 * This function can have two forms:
 * (newState: T) => T that receives 1 argument: the new state dispatched by setState, and returns the final state, or
 * (newState: T, dispatch) => void that receives 2 arguments: the new state dispatched by setState, and a function dispatch that will actually run the state update. It returns nothing.
initialState: The initial state value
 */

import {useCallback, useState, useRef} from 'react';

export default function useMediatedState(mediator, initialState) {
  const mediatorRef = useRef(mediator);

  const [state, setState] = useState(initialState);

  const setMediatedState = useCallback((nextStateOrUpdater) => {
    setState((prevState) => {
      const nextState = typeof nextStateOrUpdater === 'function'
        ? nextStateOrUpdater(prevState) 
        : nextStateOrUpdater;

      const currentMediator = mediatorRef.current;
      // has 2 value, newState and dispatch
      if (currentMediator.length >= 2) {
        let hasDispatched = false;
        let dispatchState = prevState;
        const dispatch = (valueOrUpdater) => {
          const baseState = hasDispatched ? dispatchState : prevState;
          hasDispatched = true;
          dispatchState = typeof valueOrUpdater === 'function' 
            ? valueOrUpdater(baseState)
            : valueOrUpdater;
        }
        currentMediator(nextState, dispatch);;
        return hasDispatched ? dispatchState : prevState;
      }
      // for mediator only contains 1 argument
      return currentMediator(nextState);
    });
  }, []);

  return [state, setMediatedState];
}