import { createContext, useContext, useRef } from 'react'
import { create } from 'zustand'
import { useShallow } from 'zustand/shallow'
import type {
  SplitSubvariantProps,
  SplitSubvariantProviderProps,
  SplitSubvariantState,
  SplitSubvariantStore,
} from './types.js'

const SplitSubvariantStoreContext = createContext<SplitSubvariantStore | null>(
  null
)

const shouldRecreateStore = (
  store: SplitSubvariantStore,
  props: SplitSubvariantProps
) => {
  const { state } = store.getState()
  return state !== props.state
}

export function SplitSubvariantStoreProvider({
  children,
  ...props
}: SplitSubvariantProviderProps) {
  const storeRef = useRef<SplitSubvariantStore>(null)
  if (!storeRef.current || shouldRecreateStore(storeRef.current, props)) {
    storeRef.current = createSplitSubvariantStore(props)
  }
  return (
    <SplitSubvariantStoreContext.Provider value={storeRef.current}>
      {children}
    </SplitSubvariantStoreContext.Provider>
  )
}

function useSplitSubvariantStoreContext() {
  const useStore = useContext(SplitSubvariantStoreContext)
  if (!useStore) {
    throw new Error(
      `You forgot to wrap your component in <${SplitSubvariantStoreProvider.name}>.`
    )
  }
  return useStore
}

export function useSplitSubvariantStore<T>(
  selector: (state: SplitSubvariantState) => T
): T {
  const useStore = useSplitSubvariantStoreContext()
  return useStore(useShallow(selector))
}

const createSplitSubvariantStore = ({ state }: SplitSubvariantProps) =>
  create<SplitSubvariantState>((set) => ({
    state,
    setState(state) {
      set(() => ({
        state,
      }))
    },
  }))