import { useCallback, useRef, useSyncExternalStore } from "react";
import type { AppStore } from "@editor/store";
import type { AppState } from "@editor/state";

export type AppStateSelector<TSelection> = (state: AppState) => TSelection;
export type AppStateEquality<TSelection> = (a: TSelection, b: TSelection) => boolean;

type AppStateSelectorSnapshot<TSelection> = {
  getSnapshot(state: AppState): TSelection;
  updateSnapshot(state: AppState): boolean;
};

type AppStateSelectorRef<TSelection> = {
  selector: AppStateSelector<TSelection>;
  isEqual: AppStateEquality<TSelection>;
  snapshot: AppStateSelectorSnapshot<TSelection>;
};

const selectAppState = (state: AppState) => state;

export function useAppState(store: AppStore): AppState;
export function useAppState<TSelection>(
  store: AppStore,
  selector: AppStateSelector<TSelection>,
  isEqual?: AppStateEquality<TSelection>,
): TSelection;
export function useAppState<TSelection>(
  store: AppStore,
  selector: AppStateSelector<TSelection> = selectAppState as AppStateSelector<TSelection>,
  isEqual: AppStateEquality<TSelection> = Object.is,
): TSelection {
  const selectorRef = useRef<AppStateSelectorRef<TSelection> | undefined>(undefined);

  if (!selectorRef.current || selectorRef.current.selector !== selector || selectorRef.current.isEqual !== isEqual) {
    selectorRef.current = {
      selector,
      isEqual,
      snapshot: createAppStateSelectorSnapshot(selector, isEqual),
    };
  }

  const getSnapshot = useCallback(() => selectorRef.current!.snapshot.getSnapshot(store.getState()), [store]);
  const subscribe = useCallback(
    (onStoreChange: () => void) =>
      store.subscribe((state) => {
        if (selectorRef.current!.snapshot.updateSnapshot(state)) onStoreChange();
      }),
    [store],
  );

  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
}

export function createAppStateSelectorSnapshot<TSelection>(
  selector: AppStateSelector<TSelection>,
  isEqual: AppStateEquality<TSelection> = Object.is,
): AppStateSelectorSnapshot<TSelection> {
  let stateSnapshot: AppState | undefined;
  let selectedSnapshot: TSelection | undefined;
  let hasSnapshot = false;

  const setSnapshot = (state: AppState, selected: TSelection) => {
    stateSnapshot = state;
    selectedSnapshot = selected;
    hasSnapshot = true;
  };

  return {
    getSnapshot(state) {
      if (hasSnapshot && stateSnapshot === state) return selectedSnapshot as TSelection;

      const selected = selector(state);
      if (hasSnapshot && isEqual(selectedSnapshot as TSelection, selected)) {
        stateSnapshot = state;
        return selectedSnapshot as TSelection;
      }

      setSnapshot(state, selected);
      return selected;
    },
    updateSnapshot(state) {
      const selected = selector(state);
      if (hasSnapshot && isEqual(selectedSnapshot as TSelection, selected)) {
        stateSnapshot = state;
        return false;
      }

      setSnapshot(state, selected);
      return true;
    },
  };
}

export function shallowEqual<TObject extends object>(a: TObject, b: TObject): boolean {
  if (Object.is(a, b)) return true;

  const aKeys = Object.keys(a) as Array<keyof TObject>;
  const bKeys = Object.keys(b) as Array<keyof TObject>;
  if (aKeys.length !== bKeys.length) return false;

  return aKeys.every((key) => Object.prototype.hasOwnProperty.call(b, key) && Object.is(a[key], b[key]));
}