resolutionflow/frontend/src/store/treeEditorStore.ts

import { create } from 'zustand'
import { temporal } from 'zundo'
import { immer } from 'zustand/middleware/immer'
import type { Tree, TreeStructure, TreeCreate, TreeUpdate, NodeType } from '@/types'

// Validation error interface
export interface ValidationError {
  nodeId?: string
  field?: string
  message: string
  severity: 'error' | 'warning'
}

// Draft storage key
const DRAFT_STORAGE_KEY = 'tree-editor-draft'

// Helper to generate unique IDs
const generateId = () => crypto.randomUUID()

// Helper to find a node in the tree structure
const findNodeInTree = (
  nodeId: string,
  structure: TreeStructure | null
): TreeStructure | null => {
  if (!structure) return null
  if (structure.id === nodeId) return structure
  if (structure.children) {
    for (const child of structure.children) {
      const found = findNodeInTree(nodeId, child)
      if (found) return found
    }
  }
  return null
}

// Helper to find parent of a node
const findParentNode = (
  nodeId: string,
  structure: TreeStructure | null,
  parent: TreeStructure | null = null
): TreeStructure | null => {
  if (!structure) return null
  if (structure.id === nodeId) return parent
  if (structure.children) {
    for (const child of structure.children) {
      const found = findParentNode(nodeId, child, structure)
      if (found) return found
    }
  }
  return null
}

// Helper to get all node IDs
const getAllNodeIds = (structure: TreeStructure | null): string[] => {
  if (!structure) return []
  const ids = [structure.id]
  if (structure.children) {
    for (const child of structure.children) {
      ids.push(...getAllNodeIds(child))
    }
  }
  return ids
}

// Helper to deep clone a node
const deepCloneNode = (node: TreeStructure): TreeStructure => {
  const clone: TreeStructure = { ...node, id: generateId() }

  // Update title/question to indicate it's a copy
  if (clone.question) {
    clone.question = `${clone.question} (Copy)`
  } else if (clone.title) {
    clone.title = `${clone.title} (Copy)`
  }

  // Clone options with new IDs
  if (clone.options) {
    clone.options = clone.options.map(opt => ({
      ...opt,
      id: generateId(),
      next_node_id: '' // Clear references - user must reassign
    }))
  }

  // Clear next_node_id - user must reassign
  if (clone.next_node_id) {
    clone.next_node_id = ''
  }

  // Clone children recursively
  if (clone.children) {
    clone.children = clone.children.map(child => deepCloneNode(child))
  }

  return clone
}

interface TreeEditorState {
  // Tree data
  treeId: string | null // null for new tree
  name: string
  description: string
  category: string
  treeStructure: TreeStructure | null
  originalTree: Tree | null // For comparison in edit mode

  // UI state
  selectedNodeId: string | null
  isDirty: boolean
  isLoading: boolean
  isSaving: boolean
  validationErrors: ValidationError[]

  // Auto-save state
  lastSavedAt: Date | null
  draftSavedAt: Date | null
  hasDraft: boolean

  // Actions - Initialization
  initNewTree: () => void
  loadTree: (tree: Tree) => void
  loadDraft: () => boolean
  discardDraft: () => void
  reset: () => void

  // Actions - Metadata
  setName: (name: string) => void
  setDescription: (description: string) => void
  setCategory: (category: string) => void

  // Actions - Node CRUD
  addNode: (parentId: string | null, type: NodeType, insertIndex?: number) => string
  updateNode: (nodeId: string, updates: Partial<TreeStructure>) => void
  deleteNode: (nodeId: string) => void
  duplicateNode: (nodeId: string) => string | null

  // Actions - Node ordering
  reorderNodes: (parentId: string, fromIndex: number, toIndex: number) => void
  reorderOptions: (nodeId: string, fromIndex: number, toIndex: number) => void

  // Actions - Selection
  selectNode: (nodeId: string | null) => void

  // Actions - Validation
  validate: () => ValidationError[]
  clearValidation: () => void

  // Actions - Save/Draft
  autoSaveDraft: () => void
  markSaved: () => void
  getTreeForSave: () => TreeCreate | TreeUpdate

  // Actions - State
  setLoading: (loading: boolean) => void
  setSaving: (saving: boolean) => void

  // Helpers
  findNode: (nodeId: string) => TreeStructure | null
  getAllNodeIds: () => string[]
  getAvailableTargetNodes: (excludeNodeId?: string) => Array<{ id: string; label: string; type: NodeType }>
}

// Create store with immer and temporal (undo/redo) middleware
export const useTreeEditorStore = create<TreeEditorState>()(
  temporal(
    immer((set, get) => ({
      // Initial state
      treeId: null,
      name: '',
      description: '',
      category: '',
      treeStructure: null,
      originalTree: null,
      selectedNodeId: null,
      isDirty: false,
      isLoading: false,
      isSaving: false,
      validationErrors: [],
      lastSavedAt: null,
      draftSavedAt: null,
      hasDraft: false,

      // Check for existing draft on init
      initNewTree: () => {
        const hasDraft = localStorage.getItem(DRAFT_STORAGE_KEY) !== null
        set((state) => {
          state.treeId = null
          state.name = ''
          state.description = ''
          state.category = ''
          state.treeStructure = {
            id: 'root',
            type: 'decision',
            question: '',
            options: [],
            children: []
          }
          state.originalTree = null
          state.selectedNodeId = 'root'
          state.isDirty = false
          state.isLoading = false
          state.isSaving = false
          state.validationErrors = []
          state.lastSavedAt = null
          state.draftSavedAt = null
          state.hasDraft = hasDraft
        })
      },

      loadTree: (tree: Tree) => {
        set((state) => {
          state.treeId = tree.id
          state.name = tree.name
          state.description = tree.description || ''
          state.category = tree.category || ''
          state.treeStructure = tree.tree_structure
          state.originalTree = tree
          state.selectedNodeId = tree.tree_structure?.id || null
          state.isDirty = false
          state.isLoading = false
          state.validationErrors = []
          state.lastSavedAt = new Date()
          state.draftSavedAt = null
          state.hasDraft = false
        })
      },

      loadDraft: () => {
        const draftJson = localStorage.getItem(DRAFT_STORAGE_KEY)
        if (!draftJson) return false

        try {
          const draft = JSON.parse(draftJson)
          set((state) => {
            state.treeId = draft.treeId || null
            state.name = draft.name || ''
            state.description = draft.description || ''
            state.category = draft.category || ''
            state.treeStructure = draft.treeStructure || null
            state.isDirty = true
            state.draftSavedAt = draft.savedAt ? new Date(draft.savedAt) : null
            state.hasDraft = false
          })
          return true
        } catch {
          localStorage.removeItem(DRAFT_STORAGE_KEY)
          return false
        }
      },

      discardDraft: () => {
        localStorage.removeItem(DRAFT_STORAGE_KEY)
        set((state) => {
          state.hasDraft = false
        })
      },

      reset: () => {
        set((state) => {
          state.treeId = null
          state.name = ''
          state.description = ''
          state.category = ''
          state.treeStructure = null
          state.originalTree = null
          state.selectedNodeId = null
          state.isDirty = false
          state.isLoading = false
          state.isSaving = false
          state.validationErrors = []
          state.lastSavedAt = null
          state.draftSavedAt = null
          state.hasDraft = false
        })
      },

      // Metadata actions
      setName: (name: string) => {
        set((state) => {
          state.name = name
          state.isDirty = true
        })
        get().autoSaveDraft()
      },

      setDescription: (description: string) => {
        set((state) => {
          state.description = description
          state.isDirty = true
        })
        get().autoSaveDraft()
      },

      setCategory: (category: string) => {
        set((state) => {
          state.category = category
          state.isDirty = true
        })
        get().autoSaveDraft()
      },

      // Node CRUD
      addNode: (parentId: string | null, type: NodeType, insertIndex?: number) => {
        const newId = generateId()

        const newNode: TreeStructure = {
          id: newId,
          type,
          ...(type === 'decision' && {
            question: '',
            options: [],
            children: []
          }),
          ...(type === 'action' && {
            title: '',
            description: ''
          }),
          ...(type === 'solution' && {
            title: '',
            description: ''
          })
        }

        set((state) => {
          if (!parentId) {
            // Adding as root
            state.treeStructure = newNode
          } else {
            // Find parent and add to children
            const parent = findNodeInTree(parentId, state.treeStructure)
            if (parent) {
              if (!parent.children) {
                parent.children = []
              }
              if (insertIndex !== undefined && insertIndex >= 0) {
                parent.children.splice(insertIndex, 0, newNode)
              } else {
                parent.children.push(newNode)
              }
            }
          }
          state.selectedNodeId = newId
          state.isDirty = true
        })

        get().autoSaveDraft()
        return newId
      },

      updateNode: (nodeId: string, updates: Partial<TreeStructure>) => {
        set((state) => {
          const node = findNodeInTree(nodeId, state.treeStructure)
          if (node) {
            Object.assign(node, updates)
            state.isDirty = true
          }
        })
        get().autoSaveDraft()
      },

      deleteNode: (nodeId: string) => {
        if (nodeId === 'root') {
          // Don't allow deleting root, just clear it
          set((state) => {
            if (state.treeStructure) {
              state.treeStructure.question = ''
              state.treeStructure.options = []
              state.treeStructure.children = []
            }
            state.isDirty = true
          })
          get().autoSaveDraft()
          return
        }

        set((state) => {
          const parent = findParentNode(nodeId, state.treeStructure)
          if (parent && parent.children) {
            const index = parent.children.findIndex(c => c.id === nodeId)
            if (index !== -1) {
              parent.children.splice(index, 1)
            }
          }

          // Clear selection if deleted node was selected
          if (state.selectedNodeId === nodeId) {
            state.selectedNodeId = parent?.id || 'root'
          }

          state.isDirty = true
        })
        get().autoSaveDraft()
      },

      duplicateNode: (nodeId: string) => {
        const state = get()
        const node = findNodeInTree(nodeId, state.treeStructure)
        if (!node) return null

        const clonedNode = deepCloneNode(node)

        // Find parent and add cloned node as sibling
        const parent = findParentNode(nodeId, state.treeStructure)

        set((s) => {
          if (parent && parent.children) {
            const index = parent.children.findIndex(c => c.id === nodeId)
            parent.children.splice(index + 1, 0, clonedNode)
          } else if (nodeId === 'root') {
            // Can't duplicate root - just select it
            return
          }
          s.selectedNodeId = clonedNode.id
          s.isDirty = true
        })

        get().autoSaveDraft()
        return clonedNode.id
      },

      // Reordering
      reorderNodes: (parentId: string, fromIndex: number, toIndex: number) => {
        set((state) => {
          const parent = findNodeInTree(parentId, state.treeStructure)
          if (parent && parent.children && parent.children.length > 0) {
            const [moved] = parent.children.splice(fromIndex, 1)
            parent.children.splice(toIndex, 0, moved)
            state.isDirty = true
          }
        })
        get().autoSaveDraft()
      },

      reorderOptions: (nodeId: string, fromIndex: number, toIndex: number) => {
        set((state) => {
          const node = findNodeInTree(nodeId, state.treeStructure)
          if (node && node.options && node.options.length > 0) {
            const [moved] = node.options.splice(fromIndex, 1)
            node.options.splice(toIndex, 0, moved)
            state.isDirty = true
          }
        })
        get().autoSaveDraft()
      },

      // Selection
      selectNode: (nodeId: string | null) => {
        set((state) => {
          state.selectedNodeId = nodeId
        })
      },

      // Validation
      validate: () => {
        const state = get()
        const errors: ValidationError[] = []

        // Check tree name
        if (!state.name.trim()) {
          errors.push({ message: 'Tree name is required', severity: 'error' })
        }

        // Check tree structure exists
        if (!state.treeStructure) {
          errors.push({ message: 'Tree must have at least one node', severity: 'error' })
          set((s) => { s.validationErrors = errors })
          return errors
        }

        const allNodeIds = getAllNodeIds(state.treeStructure)
        const referencedIds = new Set<string>()
        let hasSolution = false

        // Traverse and validate all nodes
        const validateNode = (node: TreeStructure) => {
          // Check type-specific required fields
          if (node.type === 'decision') {
            if (!node.question?.trim()) {
              errors.push({
                nodeId: node.id,
                field: 'question',
                message: `Decision node "${node.id}" requires a question`,
                severity: 'error'
              })
            }
            if (!node.options || node.options.length === 0) {
              errors.push({
                nodeId: node.id,
                field: 'options',
                message: `Decision node "${node.id}" requires at least one option`,
                severity: 'error'
              })
            } else {
              // Validate options
              node.options.forEach((opt, i) => {
                if (!opt.label?.trim()) {
                  errors.push({
                    nodeId: node.id,
                    field: `options[${i}].label`,
                    message: `Option ${i + 1} in "${node.id}" requires a label`,
                    severity: 'error'
                  })
                }
                if (opt.next_node_id) {
                  referencedIds.add(opt.next_node_id)
                  if (!allNodeIds.includes(opt.next_node_id)) {
                    errors.push({
                      nodeId: node.id,
                      field: `options[${i}].next_node_id`,
                      message: `Option "${opt.label}" references non-existent node "${opt.next_node_id}"`,
                      severity: 'error'
                    })
                  }
                }
              })
            }
          }

          if (node.type === 'action') {
            if (!node.title?.trim()) {
              errors.push({
                nodeId: node.id,
                field: 'title',
                message: `Action node "${node.id}" requires a title`,
                severity: 'error'
              })
            }
            if (node.next_node_id) {
              referencedIds.add(node.next_node_id)
              if (!allNodeIds.includes(node.next_node_id)) {
                errors.push({
                  nodeId: node.id,
                  field: 'next_node_id',
                  message: `Action "${node.title}" references non-existent node "${node.next_node_id}"`,
                  severity: 'error'
                })
              }
            }
          }

          if (node.type === 'solution') {
            hasSolution = true
            if (!node.title?.trim()) {
              errors.push({
                nodeId: node.id,
                field: 'title',
                message: `Solution node "${node.id}" requires a title`,
                severity: 'error'
              })
            }
          }

          // Validate children
          if (node.children) {
            node.children.forEach(child => validateNode(child))
          }
        }

        validateNode(state.treeStructure)

        // Check for at least one solution
        if (!hasSolution) {
          errors.push({
            message: 'Tree must have at least one solution (terminal) node',
            severity: 'error'
          })
        }

        // Check for orphaned nodes (not root and not referenced)
        allNodeIds.forEach(id => {
          if (id !== 'root' && !referencedIds.has(id)) {
            // Check if it's a direct child of another node (via children array)
            let isChildOfAny = false
            const checkIfChild = (node: TreeStructure) => {
              if (node.children?.some(c => c.id === id)) {
                isChildOfAny = true
              }
              node.children?.forEach(checkIfChild)
            }
            checkIfChild(state.treeStructure!)

            if (!isChildOfAny) {
              errors.push({
                nodeId: id,
                message: `Node "${id}" is orphaned (not reachable from root)`,
                severity: 'warning'
              })
            }
          }
        })

        set((s) => { s.validationErrors = errors })
        return errors
      },

      clearValidation: () => {
        set((state) => { state.validationErrors = [] })
      },

      // Auto-save draft (debounced externally, called after each change)
      autoSaveDraft: () => {
        const state = get()
        const draft = {
          treeId: state.treeId,
          name: state.name,
          description: state.description,
          category: state.category,
          treeStructure: state.treeStructure,
          savedAt: new Date().toISOString()
        }
        localStorage.setItem(DRAFT_STORAGE_KEY, JSON.stringify(draft))
        set((s) => { s.draftSavedAt = new Date() })
      },

      markSaved: () => {
        localStorage.removeItem(DRAFT_STORAGE_KEY)
        set((state) => {
          state.isDirty = false
          state.lastSavedAt = new Date()
          state.draftSavedAt = null
          state.hasDraft = false
        })
      },

      getTreeForSave: (): TreeCreate | TreeUpdate => {
        const state = get()
        return {
          name: state.name,
          description: state.description || undefined,
          category: state.category || undefined,
          tree_structure: state.treeStructure!
        }
      },

      setLoading: (loading: boolean) => {
        set((state) => { state.isLoading = loading })
      },

      setSaving: (saving: boolean) => {
        set((state) => { state.isSaving = saving })
      },

      // Helpers
      findNode: (nodeId: string) => {
        return findNodeInTree(nodeId, get().treeStructure)
      },

      getAllNodeIds: () => {
        return getAllNodeIds(get().treeStructure)
      },

      getAvailableTargetNodes: (excludeNodeId?: string) => {
        const state = get()
        const allIds = getAllNodeIds(state.treeStructure)

        return allIds
          .filter(id => id !== excludeNodeId)
          .map(id => {
            const node = findNodeInTree(id, state.treeStructure)
            let label = id
            let type: NodeType = 'decision'
            if (node) {
              type = node.type
              if (node.question) label = node.question.slice(0, 50)
              else if (node.title) label = node.title.slice(0, 50)
            }
            // Use short ID format, but 'root' stays as-is
            const shortId = id === 'root' ? 'root' : id.slice(0, 8) + '...'
            return { id, label: `${label} (${shortId})`, type }
          })
      }
    })),
    {
      // Zundo options for undo/redo
      limit: 50, // Keep last 50 states
      partialize: (state) => ({
        // Only track these fields in history
        name: state.name,
        description: state.description,
        category: state.category,
        treeStructure: state.treeStructure
      })
    }
  )
)

// Export temporal store for undo/redo access
// Use with: useStore(useTreeEditorStore.temporal, selector)
export const useTreeEditorTemporal = useTreeEditorStore.temporal

export default useTreeEditorStore