import { useEffect, useRef, useState } from 'react';
import { type Store } from './createStore';
import { shallowEqual } from '../utils/shallowEqual';

export function useStoreSelector<TState extends object, TSlice>(
  store: Store<TState>,
  selector: (state: TState) => TSlice,
  equalityFn: (a: TSlice, b: TSlice) => boolean = Object.is
): TSlice {
  const latestSelectorRef = useRef(selector);
  const latestEqualityRef = useRef(equalityFn);
  latestSelectorRef.current = selector;
  latestEqualityRef.current = equalityFn;

  const [slice, setSlice] = useState<TSlice>(() =>
    latestSelectorRef.current(store.getState())
  );

  useEffect(() => {
    function handleChange(nextState: TState, prevState: TState) {
      const nextSlice = latestSelectorRef.current(nextState);
      const prevSlice = latestSelectorRef.current(prevState);
      if (!latestEqualityRef.current(nextSlice, prevSlice)) {
        setSlice(nextSlice);
      }
    }
    const unsubscribe = store.subscribe(handleChange);
    // Sync once in case state changed between render and effect. Use the
    // functional updater form so we compare against the already-stored slice
    // rather than calling the selector twice on the same state (which returns
    // different object references for non-primitive values and causes an
    // unnecessary re-render that leads to infinite update loops with virtualizers).
    setSlice((prev) => {
      const current = latestSelectorRef.current(store.getState());
      return latestEqualityRef.current(prev, current) ? prev : current;
    });
    return unsubscribe;
  }, [store]);

  return slice;
}

export const shallow = shallowEqual;