sim/stores/workflow/middleware.ts

import { StateCreator } from 'zustand'
import { HistoryActions, HistoryEntry, WorkflowHistory } from './history-types'
import { WorkflowState, WorkflowStore } from './types'

const MAX_HISTORY_LENGTH = 20

export interface WorkflowStoreWithHistory extends WorkflowStore, HistoryActions {
  history: WorkflowHistory
}

export const withHistory = (
  config: StateCreator<WorkflowStoreWithHistory>
): StateCreator<WorkflowStoreWithHistory> => {
  return (set, get, api) => {
    const initialState = config(set, get, api)
    const initialHistoryEntry: HistoryEntry = {
      state: {
        blocks: initialState.blocks,
        edges: initialState.edges,
        loops: initialState.loops,
      },
      timestamp: Date.now(),
      action: 'Initial state',
    }

    return {
      ...initialState,
      history: {
        past: [],
        present: initialHistoryEntry,
        future: [],
      },

      canUndo: () => get().history.past.length > 0,
      canRedo: () => get().history.future.length > 0,

      undo: () => {
        const { history, ...state } = get()
        if (history.past.length === 0) return

        const previous = history.past[history.past.length - 1]
        const newPast = history.past.slice(0, history.past.length - 1)

        set({
          ...state,
          ...previous.state,
          history: {
            past: newPast,
            present: previous,
            future: [history.present, ...history.future],
          },
        })
      },

      redo: () => {
        const { history, ...state } = get()
        if (history.future.length === 0) return

        const next = history.future[0]
        const newFuture = history.future.slice(1)

        set({
          ...state,
          ...next.state,
          history: {
            past: [...history.past, history.present],
            present: next,
            future: newFuture,
          },
        })
      },

      clear: () => {
        const newState = {
          blocks: {},
          edges: [],
          loops: {},
          history: {
            past: [],
            present: {
              state: { blocks: {}, edges: [], loops: {} },
              timestamp: Date.now(),
              action: 'Clear workflow',
            },
            future: [],
          },
        }
        set(newState)
        return newState
      },

      revertToHistoryState: (index: number) => {
        const { history, ...state } = get()
        const allStates = [...history.past, history.present, ...history.future]
        const targetState = allStates[index]

        if (!targetState) return

        const newPast = allStates.slice(0, index)
        const newFuture = allStates.slice(index + 1)

        set({
          ...state,
          ...targetState.state,
          history: {
            past: newPast,
            present: targetState,
            future: newFuture,
          },
        })
      },
    }
  }
}

export const createHistoryEntry = (state: WorkflowState, action: string): HistoryEntry => ({
  state: {
    blocks: { ...state.blocks },
    edges: [...state.edges],
    loops: { ...state.loops },
  },
  timestamp: Date.now(),
  action,
})

export const pushHistory = (
  set: (
    partial:
      | Partial<WorkflowStoreWithHistory>
      | ((state: WorkflowStoreWithHistory) => Partial<WorkflowStoreWithHistory>),
    replace?: boolean
  ) => void,
  get: () => WorkflowStoreWithHistory,
  newState: WorkflowState,
  action: string
) => {
  const { history } = get()
  const newEntry = createHistoryEntry(newState, action)

  set({
    history: {
      past: [...history.past, history.present].slice(-MAX_HISTORY_LENGTH),
      present: newEntry,
      future: [],
    },
  })
}