import type {Edge, Node} from "@xyflow/react";
import type {StateCreator, StoreApi } from 'zustand/vanilla';
import type {FlowState} from "./VisProgTypes.tsx";

export type FlowSnapshot = {
  nodes: Node[];
  edges: Edge[];
}

/**
 * A reduced version of the flowState type,
 * This removes the functions that are provided by UndoRedo from the expected input type
 */
type BaseFlowState = Omit<FlowState, 'undo' | 'redo' | 'pushSnapshot' | 'beginBatchAction' | 'endBatchAction'>;


/**
 * UndoRedo is implemented as a middleware for the FlowState store,
 * this allows us to keep the undo redo logic separate from the flowState,
 * and thus from the internal editor logic
 *
 * Allows users to undo and redo actions in the visual programming editor
 *
 * @param {(set: StoreApi<FlowState>["setState"], get: () => FlowState, api: StoreApi<FlowState>) => BaseFlowState} config
 * @returns {StateCreator<FlowState>}
 * @constructor
 */
export const UndoRedo = (
  config: (
    set: StoreApi<FlowState>['setState'],
    get: () => FlowState,
    api: StoreApi<FlowState>
  ) => BaseFlowState ) : StateCreator<FlowState> => (set, get, api) => {
  let batchTimeout: number | null = null;

  /**
   * Captures the current state for
   *
   * @param {BaseFlowState} state - the current state of the editor
   * @returns {FlowSnapshot} - returns a snapshot of the current editor state
   */
  const getSnapshot = (state : BaseFlowState) : FlowSnapshot => ({
    nodes: state.nodes,
    edges: state.edges
  });

  const initialState = config(set, get, api);

  return {
    ...initialState,

    /**
     *  Adds a snapshot of the current state to the undo history
     */
    pushSnapshot: () => {
      const state = get();
      // we don't add new snapshots during an ongoing batch action
      if (!state.isBatchAction) {
        set({
          past: [...state.past, getSnapshot(state)],
          future: []
        });
      }

    },

    /**
     * Undoes the last action from the editor,
     * The state before undoing is added to the future for potential redoing
     */
    undo: () => {
      const state = get();
      if (!state.past.length) return;

      const snapshot = state.past.pop()!; // pop last snapshot
      const currentSnapshot: FlowSnapshot = getSnapshot(state);

      set({
        nodes: snapshot.nodes,
        edges: snapshot.edges,
      });

      state.future.push(currentSnapshot); // push current to redo
    },

    /**
     * redoes the last undone action,
     * The state before redoing is added to the past for potential undoing
     */
    redo: () => {
      const state = get();
      if (!state.future.length) return;

      const snapshot = state.future.pop()!; // pop last redo
      const currentSnapshot: FlowSnapshot = getSnapshot(state);

      set({
        nodes: snapshot.nodes,
        edges: snapshot.edges,
      });

      state.past.push(currentSnapshot); // push current to undo
    },

    /**
     * Begins a batched action
     *
     * An example of a batched action is dragging a node in the editor,
     * where we want the entire action of moving a node to a different position
     * to be covered by one undoable snapshot
     */
    beginBatchAction: () => {
      get().pushSnapshot();
      set({ isBatchAction: true });
      if (batchTimeout) clearTimeout(batchTimeout);
    },

    /**
     * Ends a batched action,
     * a very short timeout is used to prevent new snapshots from being added
     * until we are certain that the batch event is finished
     */
    endBatchAction: () => {
      batchTimeout = window.setTimeout(() => {
        set({ isBatchAction: false });
      }, 10);
    }
  }
}