resolutionflow/frontend/src/store/proceduralEditorStore.ts

import { create } from 'zustand'
import { temporal } from 'zundo'
import { immer } from 'zustand/middleware/immer'
import type { Tree, IntakeFormField, IntakeFieldType, ProceduralStep, ProceduralTreeStructure, TreeType } from '@/types'

const generateId = () => crypto.randomUUID()

const FIELD_TYPE_PREFIX: Record<IntakeFieldType, string> = {
  text: 'text',
  textarea: 'textarea',
  number: 'number',
  ip_address: 'ip',
  email: 'email',
  url: 'url',
  select: 'select',
  multi_select: 'multiselect',
  checkbox: 'checkbox',
  password: 'password',
}

function generateVariableName(fieldType: IntakeFieldType, existingFields: IntakeFormField[]): string {
  const prefix = FIELD_TYPE_PREFIX[fieldType] || 'field'
  const count = existingFields.filter((f) => f.field_type === fieldType).length
  return `${prefix}_${count + 1}`
}

function createDefaultStep(index: number): ProceduralStep {
  return {
    id: generateId(),
    type: 'procedure_step',
    title: `Step ${index + 1}`,
    description: '',
    content_type: 'action',
    notes_enabled: true,
  }
}

function createEndStep(): ProceduralStep {
  return {
    id: generateId(),
    type: 'procedure_end',
    title: 'Procedure Complete',
  }
}

function createSectionHeader(title: string): ProceduralStep {
  return {
    id: generateId(),
    type: 'section_header',
    title,
  }
}

function createDefaultField(index: number, existingFields: IntakeFormField[]): IntakeFormField {
  return {
    variable_name: generateVariableName('text', existingFields),
    label: `Field ${index + 1}`,
    field_type: 'text',
    required: true,
    display_order: index + 1,
  }
}

interface ProceduralEditorState {
  // Tree metadata
  treeId: string | null
  name: string
  description: string
  categoryId: string | null
  tags: string[]
  isPublic: boolean
  status: 'draft' | 'published'

  // Procedural data
  steps: ProceduralStep[]
  intakeForm: IntakeFormField[]

  // UI state
  selectedStepId: string | null
  expandedStepId: string | null
  isDirty: boolean
  isLoading: boolean
  isSaving: boolean

  // Actions - Init
  initNew: () => void
  loadTree: (tree: Tree) => void
  reset: () => void

  // Actions - Metadata
  setName: (name: string) => void
  setDescription: (description: string) => void
  setCategoryId: (categoryId: string | null) => void
  setTags: (tags: string[]) => void
  setIsPublic: (isPublic: boolean) => void
  setStatus: (status: 'draft' | 'published') => void

  // Actions - Steps
  addStep: (afterIndex?: number) => void
  addSectionHeader: (afterIndex?: number) => void
  removeStep: (stepId: string) => void
  updateStep: (stepId: string, updates: Partial<ProceduralStep>) => void
  moveStep: (fromIndex: number, toIndex: number) => void
  setSelectedStepId: (stepId: string | null) => void
  setExpandedStepId: (stepId: string | null) => void

  // Actions - Intake Form
  addField: () => void
  removeField: (index: number) => void
  updateField: (index: number, updates: Partial<IntakeFormField>) => void
  moveField: (fromIndex: number, toIndex: number) => void

  // Actions - Save
  setIsSaving: (saving: boolean) => void
  markSaved: () => void
  getTreeForSave: () => {
    name: string
    description: string
    tree_type: TreeType
    tree_structure: ProceduralTreeStructure
    intake_form: IntakeFormField[] | undefined
    category_id: string | null
    tags: string[]
    is_public: boolean
    status: 'draft' | 'published'
  }
}

export const useProceduralEditorStore = create<ProceduralEditorState>()(
  temporal(
    immer((set, get) => ({
      // Initial state
      treeId: null,
      name: '',
      description: '',
      categoryId: null,
      tags: [],
      isPublic: false,
      status: 'draft' as const,
      steps: [],
      intakeForm: [],
      selectedStepId: null,
      expandedStepId: null,
      isDirty: false,
      isLoading: false,
      isSaving: false,

      // Init
      initNew: () => {
        set((state) => {
          state.treeId = null
          state.name = ''
          state.description = ''
          state.categoryId = null
          state.tags = []
          state.isPublic = false
          state.status = 'draft'
          state.steps = [createDefaultStep(0), createEndStep()]
          state.intakeForm = []
          state.selectedStepId = null
          state.expandedStepId = null
          state.isDirty = false
          state.isLoading = false
          state.isSaving = false
        })
      },

      loadTree: (tree: Tree) => {
        const structure = tree.tree_structure as unknown as ProceduralTreeStructure
        set((state) => {
          state.treeId = tree.id
          state.name = tree.name
          state.description = tree.description || ''
          state.categoryId = tree.category_id
          state.tags = tree.tags || []
          state.isPublic = tree.is_public
          state.status = tree.status
          state.steps = structure.steps || [createDefaultStep(0), createEndStep()]
          state.intakeForm = tree.intake_form || []
          state.selectedStepId = null
          state.expandedStepId = null
          state.isDirty = false
          state.isLoading = false
          state.isSaving = false
        })
      },

      reset: () => {
        set((state) => {
          state.treeId = null
          state.name = ''
          state.description = ''
          state.categoryId = null
          state.tags = []
          state.isPublic = false
          state.status = 'draft'
          state.steps = []
          state.intakeForm = []
          state.selectedStepId = null
          state.expandedStepId = null
          state.isDirty = false
          state.isLoading = false
          state.isSaving = false
        })
      },

      // Metadata
      setName: (name) => set((state) => { state.name = name; state.isDirty = true }),
      setDescription: (description) => set((state) => { state.description = description; state.isDirty = true }),
      setCategoryId: (categoryId) => set((state) => { state.categoryId = categoryId; state.isDirty = true }),
      setTags: (tags) => set((state) => { state.tags = tags; state.isDirty = true }),
      setIsPublic: (isPublic) => set((state) => { state.isPublic = isPublic; state.isDirty = true }),
      setStatus: (status) => set((state) => { state.status = status; state.isDirty = true }),

      // Steps
      addStep: (afterIndex) => {
        set((state) => {
          // Find the insert position (before the end step)
          const endIndex = state.steps.findIndex((s) => s.type === 'procedure_end')
          const insertAt = afterIndex !== undefined
            ? Math.min(afterIndex + 1, endIndex >= 0 ? endIndex : state.steps.length)
            : (endIndex >= 0 ? endIndex : state.steps.length)

          const newStep = createDefaultStep(insertAt)
          state.steps.splice(insertAt, 0, newStep)
          state.expandedStepId = newStep.id
          state.isDirty = true
        })
      },

      addSectionHeader: (afterIndex) => {
        set((state) => {
          const endIndex = state.steps.findIndex((s) => s.type === 'procedure_end')
          const insertAt = afterIndex !== undefined
            ? Math.min(afterIndex + 1, endIndex >= 0 ? endIndex : state.steps.length)
            : (endIndex >= 0 ? endIndex : state.steps.length)

          const newHeader = createSectionHeader('New Section')
          state.steps.splice(insertAt, 0, newHeader)
          state.expandedStepId = newHeader.id
          state.isDirty = true
        })
      },

      removeStep: (stepId) => {
        set((state) => {
          const index = state.steps.findIndex((s) => s.id === stepId)
          if (index === -1) return
          // Don't remove the end step
          if (state.steps[index].type === 'procedure_end') return
          // Don't remove if it's the only procedure_step (section headers can always be removed)
          if (state.steps[index].type === 'procedure_step') {
            const stepCount = state.steps.filter((s) => s.type === 'procedure_step').length
            if (stepCount <= 1) return
          }

          state.steps.splice(index, 1)
          if (state.selectedStepId === stepId) state.selectedStepId = null
          if (state.expandedStepId === stepId) state.expandedStepId = null
          state.isDirty = true
        })
      },

      updateStep: (stepId, updates) => {
        set((state) => {
          const step = state.steps.find((s) => s.id === stepId)
          if (step) {
            Object.assign(step, updates)
            state.isDirty = true
          }
        })
      },

      moveStep: (fromIndex, toIndex) => {
        set((state) => {
          // Don't move the end step
          if (state.steps[fromIndex]?.type === 'procedure_end') return
          // Don't move past the end step
          const endIndex = state.steps.findIndex((s) => s.type === 'procedure_end')
          if (toIndex >= endIndex) return

          const [moved] = state.steps.splice(fromIndex, 1)
          state.steps.splice(toIndex, 0, moved)
          state.isDirty = true
        })
      },

      setSelectedStepId: (stepId) => set((state) => { state.selectedStepId = stepId }),
      setExpandedStepId: (stepId) => set((state) => {
        state.expandedStepId = state.expandedStepId === stepId ? null : stepId
      }),

      // Intake Form
      addField: () => {
        set((state) => {
          const newField = createDefaultField(state.intakeForm.length, state.intakeForm)
          state.intakeForm.push(newField)
          state.isDirty = true
        })
      },

      removeField: (index) => {
        set((state) => {
          if (index < 0 || index >= state.intakeForm.length) return
          state.intakeForm.splice(index, 1)
          // Reorder display_order
          state.intakeForm.forEach((f, i) => { f.display_order = i + 1 })
          state.isDirty = true
        })
      },

      updateField: (index, updates) => {
        set((state) => {
          const field = state.intakeForm[index]
          if (field) {
            // If field_type changed, auto-generate a new variable name
            if (updates.field_type && updates.field_type !== field.field_type) {
              const otherFields = state.intakeForm.filter((_, i) => i !== index)
              updates.variable_name = generateVariableName(updates.field_type, otherFields)
            }
            Object.assign(field, updates)
            state.isDirty = true
          }
        })
      },

      moveField: (fromIndex, toIndex) => {
        set((state) => {
          const [moved] = state.intakeForm.splice(fromIndex, 1)
          state.intakeForm.splice(toIndex, 0, moved)
          // Reorder display_order
          state.intakeForm.forEach((f, i) => { f.display_order = i + 1 })
          state.isDirty = true
        })
      },

      // Save
      setIsSaving: (saving) => set((state) => { state.isSaving = saving }),
      markSaved: () => set((state) => { state.isDirty = false }),

      getTreeForSave: () => {
        const state = get()
        return {
          name: state.name,
          description: state.description,
          tree_type: 'procedural' as TreeType,
          tree_structure: { steps: state.steps },
          intake_form: state.intakeForm.length > 0 ? state.intakeForm : undefined,
          category_id: state.categoryId,
          tags: state.tags,
          is_public: state.isPublic,
          status: state.status,
        }
      },
    })),
    { limit: 50 }
  )
)

export default useProceduralEditorStore