import type { TemporalState } from 'zundo'
import type { WorkflowHistoryState } from './store/workflow/history-slice'
import type { Edge, Node } from './types'
import { use, useMemo } from 'react'
import { WorkflowContext } from './context'

type WorkflowHistoryTemporalSnapshot = {
  workflowHistory: WorkflowHistoryState
}

type WorkflowHistoryTemporalStore = {
  getState: () => TemporalState<WorkflowHistoryState>
  subscribe: (listener: (state: TemporalState<WorkflowHistoryState>) => void) => () => void
}

type WorkflowHistoryStore = {
  getState: () => WorkflowHistoryState
  setState: (state: WorkflowHistoryState) => void
  subscribe: (listener: (state: WorkflowHistoryState) => void) => () => void
  temporal: WorkflowHistoryTemporalStore
}

const sanitizeWorkflowHistory = (state: WorkflowHistoryState): WorkflowHistoryState => ({
  workflowHistoryEvent: state.workflowHistoryEvent,
  workflowHistoryEventMeta: state.workflowHistoryEventMeta,
  nodes: state.nodes.map((node: Node) => ({
    ...node,
    data: {
      ...node.data,
      selected: false,
    },
  })),
  edges: state.edges.map(
    (edge: Edge) =>
      ({
        ...edge,
        selected: false,
      }) as Edge,
  ),
})

const toHistoryState = (
  state?: Partial<WorkflowHistoryTemporalSnapshot>,
): Partial<WorkflowHistoryState> => {
  return state?.workflowHistory ?? {}
}

const toTemporalState = (
  temporalState: TemporalState<WorkflowHistoryTemporalSnapshot>,
): TemporalState<WorkflowHistoryState> => ({
  pastStates: temporalState.pastStates.map(toHistoryState),
  futureStates: temporalState.futureStates.map(toHistoryState),
  undo: temporalState.undo,
  redo: temporalState.redo,
  clear: temporalState.clear,
  isTracking: temporalState.isTracking,
  pause: temporalState.pause,
  resume: temporalState.resume,
  setOnSave: (onSave) =>
    temporalState.setOnSave(
      onSave
        ? (pastState, currentState) => {
            onSave(
              toHistoryState(pastState) as WorkflowHistoryState,
              toHistoryState(currentState) as WorkflowHistoryState,
            )
          }
        : undefined,
    ),
})

export function useWorkflowHistoryStore() {
  const workflowStore = use(WorkflowContext)

  if (!workflowStore) throw new Error('Missing WorkflowContext.Provider in the tree')

  return {
    store: useMemo(
      () =>
        ({
          getState: () => workflowStore.getState().workflowHistory,
          setState: (state: WorkflowHistoryState) => {
            workflowStore.getState().setWorkflowHistory(sanitizeWorkflowHistory(state))
          },
          subscribe: (listener: (state: WorkflowHistoryState) => void) => {
            return workflowStore.subscribe((state, previousState) => {
              if (state.workflowHistory !== previousState.workflowHistory)
                listener(state.workflowHistory)
            })
          },
          temporal: {
            getState: () => toTemporalState(workflowStore.temporal.getState()),
            subscribe: (listener) =>
              workflowStore.temporal.subscribe((state) => {
                listener(toTemporalState(state))
              }),
          },
        }) satisfies WorkflowHistoryStore,
      [workflowStore],
    ),
  }
}
