pro/src/store/index.ts

import { createStore } from 'jotai'
import React, { useContext, createContext } from 'react'

export type Store = ReturnType<typeof createStore>

export const InternalPageContext = createContext<Store | undefined>(
  undefined,
)

export const useInternalStore = (): Store => {
  const store = useContext(InternalPageContext)
  if (!store) {
    throw new Error(
      `Unable to find internal Page store, Did you wrap the component within PageStoreProvider?`,
    )
  }
  return store
}

export const usePageStoreOptions = () => ({
  store: useInternalStore(),
})

export const pageStore = createStore()

export const PageStoreProvider = ({ children }: React.PropsWithChildren) => {
  const internalStoreRef = React.useRef<Store>()

  if (!internalStoreRef.current) {
    internalStoreRef.current = pageStore
  }

  return React.createElement(InternalPageContext.Provider, {
    value: internalStoreRef.current
  }, children)
}