import {
  Dispatch,
  SetStateAction,
  useCallback,
  useMemo,
  useRef,
  useState,
} from 'react';

const getCurrentState = <T>(historyAndIndex: {
  history: T[];
  index: number;
  temp: T | undefined;
}): T => {
  return historyAndIndex.temp === undefined
    ? historyAndIndex.history[historyAndIndex.index]
    : historyAndIndex.temp;
};

function defaultIsPushableStateChange<T>(prev: T, next: T) {
  return prev !== next;
}

export default function useStateHistory<T>(
  initialState: T | (() => T),
  isPushableStateChange: (
    prev: T,
    next: T
  ) => boolean = defaultIsPushableStateChange,
  onStateChange?: (state: T) => void
) {
  const [historyAndIndexWithUseState, setHistoryAndIndexWithUseState] =
    useState<{
      history: T[];
      index: number;
      temp: T | undefined;
    }>({
      history:
        initialState instanceof Function ? [initialState()] : [initialState],
      index: 0,
      temp: undefined,
    });

  const historyAndIndexWithRef = useRef<{
    history: T[];
    index: number;
    temp: T | undefined;
  }>({
    history:
      initialState instanceof Function ? [initialState()] : [initialState],
    index: 0,
    temp: undefined,
  });

  const setHistoryAndIndex: typeof setHistoryAndIndexWithUseState = useCallback(
    (update) => {
      if (!onStateChange) {
        setHistoryAndIndexWithUseState(update);
      } else {
        if (update instanceof Function) {
          historyAndIndexWithRef.current = update(
            historyAndIndexWithRef.current
          );
        } else {
          historyAndIndexWithRef.current = update;
        }
        onStateChange(getCurrentState(historyAndIndexWithRef.current));
      }
    },
    [onStateChange]
  );

  const historyAndIndex = useMemo(() => {
    return !onStateChange
      ? historyAndIndexWithUseState
      : historyAndIndexWithRef.current;
  }, [
    !onStateChange,
    historyAndIndexWithUseState,
    historyAndIndexWithRef.current,
  ]);

  const lockUndoStatesRef = useRef<boolean>(false);

  const setLockUndoStates = useCallback((lock: boolean) => {
    lockUndoStatesRef.current = lock;
  }, []);

  const undo = useCallback(() => {
    if (lockUndoStatesRef.current) return;
    setHistoryAndIndex((prev) =>
      prev.temp
        ? {
            ...prev,
            temp: undefined,
          }
        : {
            ...prev,
            index: Math.max(0, prev.index - 1),
          }
    );
  }, []);

  const redo = useCallback(() => {
    if (lockUndoStatesRef.current) return;
    setHistoryAndIndex((prev) => ({
      ...prev,
      index: Math.min(prev.history.length - 1, prev.index + 1),
      temp: undefined,
    }));
  }, []);

  const canUndo = useMemo(() => {
    return historyAndIndex.index > 0;
  }, [historyAndIndex.index]);

  const canRedo = useMemo(() => {
    return historyAndIndex.index < historyAndIndex.history.length - 1;
  }, [historyAndIndex.index, historyAndIndex.history.length]);

  const promoteTimeoutRef = useRef<NodeJS.Timeout | null>(null);

  const promoteTempState = useCallback(() => {
    if (lockUndoStatesRef.current) return;
    setHistoryAndIndex((prev) => {
      if (prev.temp === undefined) return prev;
      const pushable = isPushableStateChange(
        prev.history[prev.index],
        prev.temp
      );
      if (pushable) {
        return {
          ...prev,
          history: [...prev.history.slice(0, prev.index + 1), prev.temp],
          index: prev.index + 1,
          temp: undefined,
        };
      } else {
        return {
          ...prev,
          history: [
            ...prev.history.slice(0, prev.index),
            prev.temp,
            ...prev.history.slice(prev.index + 1),
          ],
          index: prev.index,
          temp: undefined,
        };
      }
    });
  }, [isPushableStateChange]);

  const setState = useCallback(
    (newState: T | ((prevState: T) => T)) => {
      setHistoryAndIndex((prev) => {
        const newValue =
          newState instanceof Function
            ? newState(getCurrentState(prev))
            : newState;
        return {
          ...prev,
          temp: newValue,
        };
      });
      if (promoteTimeoutRef.current) {
        clearTimeout(promoteTimeoutRef.current);
      }
      promoteTimeoutRef.current = setTimeout(promoteTempState, 150);
    },
    [promoteTempState]
  );

  const clearHistory = useCallback((state?: T) => {
    setHistoryAndIndex((prev) => ({
      history: state !== undefined ? [state] : [prev.history[prev.index]],
      index: 0,
      temp: undefined,
    }));
  }, []);

  // note: for victors only
  const rewriteHistory = useCallback((transformState: (state: T) => T) => {
    setHistoryAndIndex((prev) => ({
      ...prev,
      history: prev.history.map(transformState),
      temp: prev.temp !== undefined ? transformState(prev.temp) : undefined,
    }));
  }, []);

  const state = getCurrentState(historyAndIndex);

  const handleKeyboardEvent = useCallback(
    (e: KeyboardEvent) => {
      if (e.key === 'z' && (e.ctrlKey || e.metaKey)) {
        e.preventDefault();
        if (e.shiftKey && canRedo) {
          redo();
        } else if (!e.shiftKey && canUndo) {
          undo();
        }
      } else if (e.key === 'y' && (e.ctrlKey || e.metaKey)) {
        e.preventDefault();
        if (canRedo) {
          redo();
        }
      }
    },
    [canRedo, canUndo, redo, undo]
  );

  return useMemo(
    () => ({
      state,
      undo,
      redo,
      canUndo,
      canRedo,
      setState,
      clearHistory,
      rewriteHistory,
      setLockUndoStates,
      lockUndoStatesRef,
      handleKeyboardEvent,
    }),
    [
      state,
      undo,
      redo,
      canUndo,
      canRedo,
      setState,
      clearHistory,
      rewriteHistory,
      setLockUndoStates,
      lockUndoStatesRef,
      handleKeyboardEvent,
    ]
  );
}

export const useSynchronizingStateHistory = <T>(
  externalState: T,
  setExternalState: Dispatch<SetStateAction<T>>,
  isPushableStateChange: (
    prev: T,
    next: T
  ) => boolean = defaultIsPushableStateChange
) => {
  const initialStateRef = useRef(externalState);

  const historyOutput = useStateHistory(
    initialStateRef.current,
    isPushableStateChange,
    setExternalState // we are not going to let history changes trigger the react lifecycle. the emitted `state` will always be externalState.
  );

  return {
    ...historyOutput,
    state: externalState,
  };
};
