import { ADDON_STATE_CHANGED, ADDON_STATE_SET } from '@storybook/core-events';

import {
  HooksContext,
  applyHooks,
  useMemo,
  useCallback,
  useRef,
  useState,
  useReducer,
  useEffect,
  useChannel,
  useStoryContext,
  useParameter,
} from '@storybook/addons';

export {
  HooksContext,
  applyHooks,
  useMemo,
  useCallback,
  useRef,
  useState,
  useReducer,
  useEffect,
  useChannel,
  useStoryContext,
  useParameter,
};

export function useAddonState<S>(addonId: string, defaultState?: S): [S, (s: S) => void] {
  const [state, setState] = useState<S>(defaultState);
  const emit = useChannel(
    {
      [`${ADDON_STATE_CHANGED}-manager-${addonId}`]: (s: S) => {
        setState(s);
      },
      [`${ADDON_STATE_SET}-manager-${addonId}`]: (s: S) => {
        setState(s);
      },
    },
    [addonId]
  );

  useEffect(() => {
    // init
    if (defaultState !== undefined) {
      emit(`${ADDON_STATE_SET}-client-${addonId}`, defaultState);
    }
  }, []);

  return [
    state,
    s => {
      setState(s);
      emit(`${ADDON_STATE_CHANGED}-client-${addonId}`, s);
    },
  ];
}

import { ADDON_STATE_CHANGED, ADDON_STATE_SET } from '@storybook/core-events';

import {
  HooksContext,
  applyHooks,
  useMemo,
  useCallback,
  useRef,
  useState,
  useReducer,
  useEffect,
  useChannel,
  useStoryContext,
  useParameter,
} from '@storybook/addons';

export {
  HooksContext,
  applyHooks,
  useMemo,
  useCallback,
  useRef,
  useState,
  useReducer,
  useEffect,
  useChannel,
  useStoryContext,
  useParameter,
};

// We keep this store, because when stories are edited by the user, and HMR, state is lost.
// This allows us to restore instantly.
const addonStateCache: Record<string, any> = {};

export function useAddonState<S>(addonId: string, defaultState?: S): [S, (s: S) => void] {
  const restoredState = addonStateCache[addonId] as S;

  const [state, setState] = useState<S>(
    typeof restoredState === 'undefined' ? defaultState : restoredState
  );

  addonStateCache[addonId] = state;

  const allListeners = useMemo(
    () => ({
      [`${ADDON_STATE_CHANGED}-manager-${addonId}`]: (s: S) => setState(s),
      [`${ADDON_STATE_SET}-manager-${addonId}`]: (s: S) => setState(s),
    }),
    [addonId]
  );

  const emit = useChannel(allListeners, [addonId]);

  useEffect(() => {
    // init
    if (defaultState !== undefined) {
      emit(`${ADDON_STATE_SET}-client-${addonId}`, defaultState);
    }
  }, []);

  return [
    state,
    s => {
      setState(s);
      emit(`${ADDON_STATE_CHANGED}-client-${addonId}`, s);
    },
  ];
}

ADD memoization for eventHandlers object & ADD a addonState cache for HMR compat
ADD memoization for eventHandlers object & ADD a addonState cache for HMR compat
