import React, {
  useCallback,
  useContext,
  useMemo,
  useRef,
  useState,
} from "react";

/** This is a performance optimization. It allows using state throughout the
 *  component tree while minimizing rerenders
 *
 */
export const createStateContext = <T extends unknown>(initialValue: T) => {
  const stateContext = React.createContext<T>(initialValue);
  const staticContext = React.createContext<{
    setState: React.Dispatch<React.SetStateAction<T>>;
    ref: { current: T };
  }>({
    setState: () => {},
    ref: { current: initialValue },
  });

  return {
    // useState is used when you want to have the state value and you need to
    // rerender your component every time the state value changes
    // eg. when the state changes how a component renders
    useState: () => useContext(stateContext),
    // useRef is used when you want to have the state value and don't want to
    // rerender your component when the state value changes
    // eg. when the state is used in a callback
    useRef: () => useContext(staticContext).ref as { readonly current: T },
    // useSetter to update the state value. This value should never change so it
    // won't cause rerenders.
    useSetter: () => {
      const { setState, ref } = useContext(staticContext);
      return useCallback(
        (action: React.SetStateAction<T>) => {
          const newValue =
            typeof action === "function"
              ? (action as (val: T) => T)(ref.current)
              : action;
          ref.current = newValue;
          setState(newValue);
        },
        [setState, ref]
      );
    },
    Provider: ({ children }: { children: React.ReactNode }) => {
      const [state, setState] = useState(initialValue);
      const ref = useRef(initialValue);
      const staticValue = useMemo(() => ({ setState, ref }), []);

      return (
        <stateContext.Provider value={state}>
          <staticContext.Provider value={staticValue}>
            {children}
          </staticContext.Provider>
        </stateContext.Provider>
      );
    },
  };
};