simstudioai--sim/apps/sim/stores/workflow-diff/store.ts

import { createLogger } from '@sim/logger'
import { create } from 'zustand'
import { devtools } from 'zustand/middleware'
import { stripWorkflowDiffMarkers, WorkflowDiffEngine } from '@/lib/workflows/diff'
import { enqueueReplaceWorkflowState } from '@/lib/workflows/operations/socket-operations'
import { validateWorkflowState } from '@/lib/workflows/sanitization/validation'
import { Serializer } from '@/serializer'
import { useWorkflowRegistry } from '../workflows/registry/store'
import { mergeSubblockState } from '../workflows/utils'
import { useWorkflowStore } from '../workflows/workflow/store'
import { deriveDiffFlags, type WorkflowDiffActions, type WorkflowDiffState } from './types'
import {
  applyWorkflowStateToStores,
  captureBaselineSnapshot,
  cloneWorkflowState,
  createBatchedUpdater,
  getLatestUserMessageId,
  persistWorkflowStateToServer,
  WORKFLOW_DIFF_SETTLED_EVENT,
} from './utils'

const logger = createLogger('WorkflowDiffStore')
const diffEngine = new WorkflowDiffEngine()

/**
 * Canonical state patch that clears the diff overlay back to `none`: the
 * none-derived flags plus a wipe of all diff payload fields.
 */
export const RESET_DIFF_STATE = {
  ...deriveDiffFlags('none'),
  baselineWorkflow: null,
  baselineWorkflowId: null,
  diffAnalysis: null,
  diffMetadata: null,
  diffError: null,
  _triggerMessageId: null,
} as const

/**
 * Detects when a diff contains no meaningful changes.
 */
function isEmptyDiffAnalysis(
  diffAnalysis?: {
    new_blocks?: string[]
    edited_blocks?: string[]
    deleted_blocks?: string[]
    field_diffs?: Record<string, { changed_fields: string[] }>
    edge_diff?: { new_edges?: string[]; deleted_edges?: string[] }
  } | null
): boolean {
  if (!diffAnalysis) return false
  const hasBlockChanges =
    (diffAnalysis.new_blocks?.length || 0) > 0 ||
    (diffAnalysis.edited_blocks?.length || 0) > 0 ||
    (diffAnalysis.deleted_blocks?.length || 0) > 0
  const hasEdgeChanges =
    (diffAnalysis.edge_diff?.new_edges?.length || 0) > 0 ||
    (diffAnalysis.edge_diff?.deleted_edges?.length || 0) > 0
  const hasFieldChanges = Object.values(diffAnalysis.field_diffs || {}).some(
    (diff) => (diff?.changed_fields?.length || 0) > 0
  )
  return !hasBlockChanges && !hasEdgeChanges && !hasFieldChanges
}

function notifyDiffSettled(workflowId: string | null | undefined): void {
  if (!workflowId || typeof window === 'undefined') return
  window.dispatchEvent(new CustomEvent(WORKFLOW_DIFF_SETTLED_EVENT, { detail: { workflowId } }))
}

export const useWorkflowDiffStore = create<WorkflowDiffState & WorkflowDiffActions>()(
  devtools(
    (set, get) => {
      const batchedUpdate = createBatchedUpdater(set)

      return {
        ...deriveDiffFlags('none'),
        baselineWorkflow: null,
        baselineWorkflowId: null,
        diffAnalysis: null,
        diffMetadata: null,
        diffError: null,
        pendingExternalUpdates: {},
        remoteUpdateVersions: {},
        reconcilingWorkflows: {},
        reconciliationErrors: {},
        _triggerMessageId: null,
        _batchedStateUpdate: batchedUpdate,

        setProposedChanges: async (proposedState, diffAnalysis, options) => {
          const activeWorkflowId = useWorkflowRegistry.getState().activeWorkflowId
          if (!activeWorkflowId) {
            logger.error('Cannot apply diff without an active workflow')
            throw new Error('No active workflow found')
          }

          // Capture baseline if needed (synchronous, fast)
          let baselineWorkflow = get().baselineWorkflow
          let baselineWorkflowId = get().baselineWorkflowId
          let capturedBaseline = false

          if (
            options?.baselineWorkflow &&
            (!baselineWorkflow || baselineWorkflowId !== activeWorkflowId)
          ) {
            baselineWorkflow = cloneWorkflowState(options.baselineWorkflow)
            baselineWorkflowId = activeWorkflowId
            capturedBaseline = true
            logger.info('Using provided baseline snapshot for diff workflow', {
              workflowId: activeWorkflowId,
              blockCount: Object.keys(baselineWorkflow.blocks || {}).length,
            })
          } else if (!baselineWorkflow || baselineWorkflowId !== activeWorkflowId) {
            baselineWorkflow = captureBaselineSnapshot(activeWorkflowId)
            baselineWorkflowId = activeWorkflowId
            capturedBaseline = true
            logger.info('Captured baseline snapshot for diff workflow', {
              workflowId: activeWorkflowId,
              blockCount: Object.keys(baselineWorkflow.blocks || {}).length,
            })
          }

          // Create diff (this is fast, just computes the diff)
          const diffResult = await diffEngine.createDiffFromWorkflowState(
            proposedState,
            diffAnalysis,
            baselineWorkflow ?? undefined
          )

          if (!diffResult.success || !diffResult.diff) {
            const errorMessage = diffResult.errors?.join(', ') || 'Failed to create diff'
            logger.error(errorMessage)
            throw new Error(errorMessage)
          }

          const diffAnalysisResult = diffResult.diff.diffAnalysis || null
          if (isEmptyDiffAnalysis(diffAnalysisResult)) {
            logger.info('No workflow diff detected; skipping diff view')
            diffEngine.clearDiff()
            batchedUpdate(RESET_DIFF_STATE)
            return
          }

          const candidateState = diffResult.diff.proposedState

          logger.info('[WorkflowDiff] Applying proposed state', {
            blockCount: Object.keys(candidateState.blocks || {}).length,
            edgeCount: candidateState.edges?.length ?? 0,
            hasLoops: !!candidateState.loops,
            hasParallels: !!candidateState.parallels,
          })

          // Validate proposed workflow using serializer round-trip
          const serializer = new Serializer()
          const serialized = serializer.serializeWorkflow(
            candidateState.blocks,
            candidateState.edges,
            candidateState.loops,
            candidateState.parallels,
            false
          )
          serializer.deserializeWorkflow(serialized)

          // Set diff status BEFORE applying state so that the very first
          // React render after the store update already sees diff mode active.
          // This prevents a flash frame where blocks render without diff colors.
          set({
            ...deriveDiffFlags('showing'),
            baselineWorkflow: baselineWorkflow,
            baselineWorkflowId,
            diffAnalysis: diffAnalysisResult,
            diffMetadata: diffResult.diff.metadata,
            diffError: null,
            _triggerMessageId: get()._triggerMessageId ?? null,
          })

          // Now apply the proposed state to stores — React will batch this
          // with the diff flag update above into a single render.
          applyWorkflowStateToStores(activeWorkflowId, candidateState)
          logger.info('[WorkflowDiff] Applied state to stores with diff flags already active')

          // Resolve trigger message ID in the background (non-blocking)
          if (capturedBaseline && !get()._triggerMessageId) {
            getLatestUserMessageId().then((msgId) => {
              if (msgId) set({ _triggerMessageId: msgId })
            })
          }

          logger.info('Workflow diff applied optimistically', {
            workflowId: activeWorkflowId,
            blocks: Object.keys(candidateState.blocks || {}).length,
            edges: candidateState.edges?.length || 0,
          })

          // When skipPersist is set, the server tool (edit_workflow) already
          // saved to DB. Both the Socket.IO broadcast and HTTP persist would
          // race with subsequent edit_workflow calls and overwrite newer state,
          // causing block IDs to thrash.
          if (!options?.skipPersist) {
            const cleanState = stripWorkflowDiffMarkers(cloneWorkflowState(candidateState))

            enqueueReplaceWorkflowState({
              workflowId: activeWorkflowId,
              state: cleanState,
            }).catch((error) => {
              logger.warn('Failed to broadcast workflow state (non-blocking)', { error })
            })

            persistWorkflowStateToServer(activeWorkflowId, candidateState)
              .then((persisted) => {
                if (!persisted) {
                  logger.warn('Failed to persist copilot edits (state already applied locally)')
                } else {
                  logger.info('Workflow diff persisted to database', {
                    workflowId: activeWorkflowId,
                  })
                }
              })
              .catch((error) => {
                logger.warn('Failed to persist workflow state (non-blocking)', { error })
              })
          }

          // Emit event for undo/redo recording
          if (!options?.skipRecording) {
            window.dispatchEvent(
              new CustomEvent('record-diff-operation', {
                detail: {
                  type: 'apply-diff',
                  baselineSnapshot: baselineWorkflow,
                  proposedState: candidateState,
                  diffAnalysis: diffResult.diff.diffAnalysis,
                },
              })
            )
          }
        },

        clearDiff: ({ restoreBaseline = true } = {}) => {
          const { baselineWorkflow, baselineWorkflowId } = get()
          const activeWorkflowId = useWorkflowRegistry.getState().activeWorkflowId

          if (
            restoreBaseline &&
            baselineWorkflow &&
            baselineWorkflowId &&
            baselineWorkflowId === activeWorkflowId
          ) {
            applyWorkflowStateToStores(baselineWorkflowId, baselineWorkflow)
          }

          diffEngine.clearDiff()

          set(RESET_DIFF_STATE)
          notifyDiffSettled(baselineWorkflowId ?? activeWorkflowId)
        },

        toggleDiffView: () => {
          const { status } = get()
          if (status === 'none') {
            logger.warn('Cannot toggle diff view without an active, ready diff')
            return
          }
          batchedUpdate(deriveDiffFlags(status === 'showing' ? 'staged' : 'showing'))
        },

        acceptChanges: async (options) => {
          const activeWorkflowId = useWorkflowRegistry.getState().activeWorkflowId
          if (!activeWorkflowId) {
            logger.error('No active workflow ID found when accepting diff')
            throw new Error('No active workflow found')
          }

          const workflowStore = useWorkflowStore.getState()
          const currentState = workflowStore.getWorkflowState()
          const mergedBlocks = mergeSubblockState(
            currentState.blocks,
            activeWorkflowId ?? undefined
          )
          const mergedState = {
            ...currentState,
            blocks: mergedBlocks,
          }
          const cleanState = stripWorkflowDiffMarkers(cloneWorkflowState(mergedState))
          const validation = validateWorkflowState(cleanState, { sanitize: true })

          if (!validation.valid) {
            const errorMessage = `Cannot apply changes: ${validation.errors.join('; ')}`
            logger.error(errorMessage)
            batchedUpdate({ diffError: errorMessage })
            throw new Error(errorMessage)
          }

          const stateToApply = {
            ...(validation.sanitizedState || cleanState),
            lastSaved: useWorkflowStore.getState().lastSaved,
          }

          // Capture state before accept for undo
          const beforeAccept = cloneWorkflowState(mergedState)
          const afterAccept = cloneWorkflowState(stateToApply)
          const diffAnalysisForUndo = get().diffAnalysis
          const baselineForUndo = get().baselineWorkflow

          // Clear diff state FIRST to prevent flash of colors
          // This must happen synchronously before applying the cleaned state
          set(RESET_DIFF_STATE)

          // Clear the diff engine
          diffEngine.clearDiff()

          // Now apply the cleaned state
          applyWorkflowStateToStores(activeWorkflowId, stateToApply)

          // Emit event for undo/redo recording (unless we're in an undo/redo operation)
          if (!options?.skipRecording) {
            window.dispatchEvent(
              new CustomEvent('record-diff-operation', {
                detail: {
                  type: 'accept-diff',
                  beforeAccept,
                  afterAccept,
                  diffAnalysis: diffAnalysisForUndo,
                  baselineSnapshot: baselineForUndo,
                },
              })
            )
          }

          notifyDiffSettled(activeWorkflowId)
        },

        rejectChanges: async (options) => {
          const { baselineWorkflow, baselineWorkflowId, diffAnalysis } = get()
          const activeWorkflowId = useWorkflowRegistry.getState().activeWorkflowId

          if (!baselineWorkflow || !baselineWorkflowId) {
            logger.warn('Reject called without baseline workflow')
            get().clearDiff({ restoreBaseline: false })
            return
          }

          if (!activeWorkflowId || activeWorkflowId !== baselineWorkflowId) {
            logger.warn('Reject called while viewing a different workflow', {
              activeWorkflowId,
              baselineWorkflowId,
            })
            get().clearDiff({ restoreBaseline: false })
            return
          }

          // Capture current state (with markers) before rejecting
          const workflowStore = useWorkflowStore.getState()
          const currentState = workflowStore.getWorkflowState()
          const mergedBlocks = mergeSubblockState(
            currentState.blocks,
            activeWorkflowId ?? undefined
          )
          const beforeReject = cloneWorkflowState({
            ...currentState,
            blocks: mergedBlocks,
          })

          get().setWorkflowReconciliationInProgress(baselineWorkflowId, true)

          // Clear diff state FIRST for instant UI feedback
          set(RESET_DIFF_STATE)

          // Clear the diff engine
          diffEngine.clearDiff()

          // Apply baseline state locally
          applyWorkflowStateToStores(baselineWorkflowId, baselineWorkflow)

          // Emit event for undo/redo recording synchronously
          if (!options?.skipRecording) {
            window.dispatchEvent(
              new CustomEvent('record-diff-operation', {
                detail: {
                  type: 'reject-diff',
                  beforeReject,
                  afterReject: baselineWorkflow,
                  diffAnalysis,
                  baselineSnapshot: baselineWorkflow,
                },
              })
            )
          }

          logger.info('Broadcasting reject to other users', {
            workflowId: activeWorkflowId,
            blockCount: Object.keys(baselineWorkflow.blocks).length,
          })

          enqueueReplaceWorkflowState({
            workflowId: activeWorkflowId,
            state: baselineWorkflow,
          }).catch((error) => {
            logger.error('Failed to broadcast reject to other users:', error)
          })

          const pendingGenerationBeforePersist =
            get().pendingExternalUpdates[baselineWorkflowId] ?? 0
          const persisted = await persistWorkflowStateToServer(baselineWorkflowId, baselineWorkflow)
          if (!persisted) {
            logger.error('Failed to persist baseline workflow state')
            if (
              (get().pendingExternalUpdates[baselineWorkflowId] ?? 0) <=
              pendingGenerationBeforePersist
            ) {
              get().clearExternalUpdatePending(baselineWorkflowId)
            }
            get().setWorkflowReconciliationInProgress(baselineWorkflowId, false)
            get().setWorkflowReconciliationError(
              baselineWorkflowId,
              'Failed to save rejected copilot changes. Refresh and try again.'
            )
            if (
              (get().pendingExternalUpdates[baselineWorkflowId] ?? 0) >
              pendingGenerationBeforePersist
            ) {
              notifyDiffSettled(baselineWorkflowId)
            }
            return
          }

          if (
            (get().pendingExternalUpdates[baselineWorkflowId] ?? 0) <=
            pendingGenerationBeforePersist
          ) {
            get().clearExternalUpdatePending(baselineWorkflowId)
          }
          get().setWorkflowReconciliationError(baselineWorkflowId, null)
          get().setWorkflowReconciliationInProgress(baselineWorkflowId, false)
          notifyDiffSettled(baselineWorkflowId)
        },

        reapplyDiffMarkers: () => {
          const { hasActiveDiff, isDiffReady, diffAnalysis } = get()
          if (!hasActiveDiff || !isDiffReady || !diffAnalysis) {
            return
          }

          const workflowStore = useWorkflowStore.getState()
          const currentBlocks = workflowStore.blocks

          // Check if any blocks need markers applied (checking the actual property, not just existence)
          const needsUpdate =
            diffAnalysis.new_blocks?.some((blockId) => {
              const block = currentBlocks[blockId]
              const blockDiffState = (block as { is_diff?: string } | undefined)?.is_diff
              return block && blockDiffState !== 'new'
            }) ||
            diffAnalysis.edited_blocks?.some((blockId) => {
              const block = currentBlocks[blockId]
              const blockDiffState = (block as { is_diff?: string } | undefined)?.is_diff
              return block && blockDiffState !== 'edited'
            })

          if (!needsUpdate) {
            return
          }

          const updatedBlocks: Record<string, any> = {}
          let hasChanges = false

          // Only clone blocks that need diff markers
          Object.entries(currentBlocks).forEach(([blockId, block]) => {
            const isNewBlock = diffAnalysis.new_blocks?.includes(blockId)
            const isEditedBlock = diffAnalysis.edited_blocks?.includes(blockId)
            const blockDiffState = (block as { is_diff?: string } | undefined)?.is_diff

            if (isNewBlock && blockDiffState !== 'new') {
              updatedBlocks[blockId] = { ...block, is_diff: 'new' }
              hasChanges = true
            } else if (isEditedBlock && blockDiffState !== 'edited') {
              updatedBlocks[blockId] = { ...block, is_diff: 'edited' }

              // Re-apply field_diffs if available
              if (diffAnalysis.field_diffs?.[blockId]) {
                updatedBlocks[blockId].field_diffs = diffAnalysis.field_diffs[blockId]

                // Clone subblocks and apply markers
                const fieldDiff = diffAnalysis.field_diffs[blockId]
                updatedBlocks[blockId].subBlocks = { ...block.subBlocks }

                fieldDiff.changed_fields.forEach((field) => {
                  if (updatedBlocks[blockId].subBlocks?.[field]) {
                    updatedBlocks[blockId].subBlocks[field] = {
                      ...updatedBlocks[blockId].subBlocks[field],
                      is_diff: 'changed',
                    }
                  }
                })
              }
              hasChanges = true
            } else {
              updatedBlocks[blockId] = block
            }
          })

          // Only update if we actually made changes
          if (hasChanges) {
            useWorkflowStore.setState({ blocks: updatedBlocks })
            logger.info('Re-applied diff markers to workflow blocks')
          }
        },

        markRemoteUpdateSeen: (workflowId) => {
          set((state) => ({
            remoteUpdateVersions: {
              ...state.remoteUpdateVersions,
              [workflowId]: (state.remoteUpdateVersions[workflowId] ?? 0) + 1,
            },
            reconciliationErrors: Object.fromEntries(
              Object.entries(state.reconciliationErrors).filter(([id]) => id !== workflowId)
            ),
          }))
        },

        markExternalUpdatePending: (workflowId) => {
          const current = get()
          set({
            pendingExternalUpdates: {
              ...current.pendingExternalUpdates,
              [workflowId]: (current.pendingExternalUpdates[workflowId] ?? 0) + 1,
            },
            remoteUpdateVersions: {
              ...current.remoteUpdateVersions,
              [workflowId]: (current.remoteUpdateVersions[workflowId] ?? 0) + 1,
            },
            reconciliationErrors: Object.fromEntries(
              Object.entries(current.reconciliationErrors).filter(([id]) => id !== workflowId)
            ),
          })
        },

        clearExternalUpdatePending: (workflowId) => {
          set((state) => {
            const { [workflowId]: _removed, ...pendingExternalUpdates } =
              state.pendingExternalUpdates
            return { pendingExternalUpdates }
          })
        },

        setWorkflowReconciliationInProgress: (workflowId, isReconciling) => {
          set((state) => {
            const { [workflowId]: _removed, ...reconcilingWorkflows } = state.reconcilingWorkflows
            return {
              reconcilingWorkflows: isReconciling
                ? { ...reconcilingWorkflows, [workflowId]: true }
                : reconcilingWorkflows,
            }
          })
        },

        setWorkflowReconciliationError: (workflowId, error) => {
          set((state) => {
            const { [workflowId]: _removed, ...reconciliationErrors } = state.reconciliationErrors
            return {
              reconciliationErrors: error
                ? { ...reconciliationErrors, [workflowId]: error }
                : reconciliationErrors,
            }
          })
        },
      }
    },
    { name: 'workflow-diff-store' }
  )
)