resolutionflow/frontend/src/components/tree-editor/TreeCanvas.tsx

import { useState, useCallback, useRef, useEffect } from 'react'
import { HelpCircle, Zap, CheckCircle, Plus, X } from 'lucide-react'
import { useTreeEditorStore, findNodeInTree } from '@/store/treeEditorStore'
import { TreeCanvasNode } from './TreeCanvasNode'
import { AnswerStubCard } from './AnswerStubCard'
import type { TreeStructure, NodeType } from '@/types'
import { cn } from '@/lib/utils'

// ─── Types ───────────────────────────────────────────────────────────────────

interface PendingLink {
  parentId: string
  optionId?: string  // For decision option linking
}

interface DragState {
  nodeId: string
  parentId: string | null
  index: number
}

// ─── Reference cleanup helper ─────────────────────────────────────────────────

/**
 * Before deleting a node, clear all inbound references to it across the tree.
 * This prevents stale next_node_id / option.next_node_id references.
 */
function clearInboundReferences(
  nodeId: string,
  treeStructure: TreeStructure,
  updateNode: (id: string, updates: Partial<TreeStructure>) => void
) {
  function walk(node: TreeStructure) {
    // Clear decision option references
    if (node.type === 'decision' && node.options) {
      const needsUpdate = node.options.some((o) => o.next_node_id === nodeId)
      if (needsUpdate) {
        updateNode(node.id, {
          options: node.options.map((o) =>
            o.next_node_id === nodeId ? { ...o, next_node_id: '' } : o
          ),
        })
      }
    }

    // Clear action next_node_id references
    if (node.type === 'action' && node.next_node_id === nodeId) {
      updateNode(node.id, { next_node_id: '' })
    }

    // Recurse
    node.children?.forEach(walk)
  }

  walk(treeStructure)
}

// ─── Add-node type picker ─────────────────────────────────────────────────────

interface AddNodePickerProps {
  onSelect: (type: NodeType) => void
  onCancel: () => void
}

function AddNodePicker({ onSelect, onCancel }: AddNodePickerProps) {
  return (
    <div className="flex items-center gap-2 rounded-xl border border-dashed border-primary/40 bg-[#14161d] px-3 py-2 shadow-xs">
      <span className="text-xs text-[#848b9b] shrink-0">Add:</span>

      <button
        type="button"
        onClick={() => onSelect('decision')}
        className={cn(
          'flex items-center gap-1 rounded-md px-2 py-1 text-xs font-sans text-xs',
          'border border-blue-500/30 bg-blue-500/10 text-blue-400 hover:bg-blue-500/20'
        )}
      >
        <HelpCircle className="h-3 w-3" />
        Decision
      </button>

      <button
        type="button"
        onClick={() => onSelect('action')}
        className={cn(
          'flex items-center gap-1 rounded-md px-2 py-1 text-xs font-sans text-xs',
          'border border-yellow-500/30 bg-yellow-500/10 text-yellow-400 hover:bg-yellow-500/20'
        )}
      >
        <Zap className="h-3 w-3" />
        Action
      </button>

      <button
        type="button"
        onClick={() => onSelect('solution')}
        className={cn(
          'flex items-center gap-1 rounded-md px-2 py-1 text-xs font-sans text-xs',
          'border border-green-500/30 bg-green-500/10 text-green-400 hover:bg-green-500/20'
        )}
      >
        <CheckCircle className="h-3 w-3" />
        Solution
      </button>

      <button
        type="button"
        onClick={onCancel}
        className="ml-1 rounded p-0.5 text-[#848b9b] hover:bg-accent"
      >
        <X className="h-3 w-3" />
      </button>
    </div>
  )
}

// ─── Add-node trigger button ──────────────────────────────────────────────────

interface AddNodeButtonProps {
  label?: string
  onClick: () => void
}

function AddNodeButton({ label = 'Add node', onClick }: AddNodeButtonProps) {
  return (
    <button
      type="button"
      onClick={onClick}
      className={cn(
        'flex items-center gap-1 rounded-lg px-3 py-1.5 text-xs font-sans text-xs',
        'border border-dashed border-[#1e2130] text-[#848b9b]',
        'hover:border-primary/40 hover:text-[#e2e5eb] hover:bg-accent/50',
        'transition-all duration-150'
      )}
    >
      <Plus className="h-3 w-3" />
      {label}
    </button>
  )
}

// ─── Add-key builder ──────────────────────────────────────────────────────────

/** Unique key for an add-target: "parentId" or "parentId:optionId" */
function addKey(parentId: string, optionId?: string) {
  return optionId ? `${parentId}:${optionId}` : parentId
}

// ─── TreeCanvas ───────────────────────────────────────────────────────────────

export function TreeCanvas() {
  const {
    treeStructure,
    addNode,
    updateNode,
    deleteNode,
    duplicateNode,
    reorderNodes,
    selectNode,
    selectedNodeId,
  } = useTreeEditorStore()

  // ── Local canvas state ──
  const [expandedNodeId, setExpandedNodeId] = useState<string | null>(null)
  const [newNodeIds, setNewNodeIds] = useState<Set<string>>(new Set())
  const [collapsedNodeIds, setCollapsedNodeIds] = useState<Set<string>>(new Set())
  const [pendingAddKey, setPendingAddKey] = useState<string | null>(null)
  const [pendingLinks, setPendingLinks] = useState<Map<string, PendingLink>>(
    new Map()
  )
  const [dragState, setDragState] = useState<DragState | null>(null)
  const [dragOverTarget, setDragOverTarget] = useState<{
    parentId: string | null
    index: number
  } | null>(null)

  // Node ref map for scroll-into-view
  const nodeRefs = useRef<Map<string, HTMLDivElement>>(new Map())

  // ── Selection sync ──
  // When selectedNodeId changes externally (e.g. ValidationSummary click),
  // auto-expand that card and scroll it into view.
  useEffect(() => {
    if (selectedNodeId && selectedNodeId !== expandedNodeId) {
      setExpandedNodeId(selectedNodeId)
      const el = nodeRefs.current.get(selectedNodeId)
      if (el) {
        el.scrollIntoView({ behavior: 'smooth', block: 'nearest' })
      }
    }
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [selectedNodeId])

  // ── Card expand/collapse ──
  const handleToggleExpand = useCallback(
    (nodeId: string) => {
      setExpandedNodeId((prev) => (prev === nodeId ? null : nodeId))
      selectNode(nodeId)
    },
    [selectNode]
  )

  // ── Save inline edits ──
  const handleSave = useCallback(
    (nodeId: string, updates: Partial<TreeStructure>) => {
      updateNode(nodeId, updates)

      // For decision nodes: strip blank options, then create answer stubs for any
      // labelled option that doesn't yet have a linked child
      if (updates.options) {
        const options = updates.options.filter((o) => o.label.trim())
        const stubsToCreate: Array<{ opt: typeof options[number]; stubId: string }> = []

        options.forEach((opt) => {
          if (!opt.next_node_id) {
            const stubId = addNode(nodeId, 'answer')
            updateNode(stubId, { title: opt.label })
            stubsToCreate.push({ opt, stubId })
          }
        })

        // Write back: filtered options + any newly assigned next_node_ids
        const updatedOptions = options.map((o) => {
          const stub = stubsToCreate.find((s) => s.opt.id === o.id)
          return stub ? { ...o, next_node_id: stub.stubId } : o
        })
        updateNode(nodeId, { options: updatedOptions })
      }

      // Resolve pending link for new nodes
      const link = pendingLinks.get(nodeId)
      if (link) {
        const parentNode = treeStructure
          ? findNodeInTree(link.parentId, treeStructure)
          : null

        if (parentNode) {
          if (link.optionId && parentNode.type === 'decision' && parentNode.options) {
            // Link the decision option to this new child node
            const updatedOptions = parentNode.options.map((o) =>
              o.id === link.optionId ? { ...o, next_node_id: nodeId } : o
            )
            updateNode(link.parentId, { options: updatedOptions })
          } else if (parentNode.type === 'action') {
            // Link the action's next node
            updateNode(link.parentId, { next_node_id: nodeId })
          }
        }

        setPendingLinks((prev) => {
          const next = new Map(prev)
          next.delete(nodeId)
          return next
        })
      }

      setNewNodeIds((prev) => {
        const next = new Set(prev)
        next.delete(nodeId)
        return next
      })
      setExpandedNodeId(null)
    },
    [pendingLinks, treeStructure, updateNode]
  )

  // ── Cancel new node ──
  const handleCancelNew = useCallback(
    (nodeId: string) => {
      deleteNode(nodeId)
      setNewNodeIds((prev) => {
        const next = new Set(prev)
        next.delete(nodeId)
        return next
      })
      setPendingLinks((prev) => {
        const next = new Map(prev)
        next.delete(nodeId)
        return next
      })
      if (expandedNodeId === nodeId) setExpandedNodeId(null)
    },
    [deleteNode, expandedNodeId]
  )

  // ── Delete node (with inbound reference cleanup) ──
  const handleDelete = useCallback(
    (nodeId: string) => {
      if (!treeStructure) return
      clearInboundReferences(nodeId, treeStructure, updateNode)
      deleteNode(nodeId)
      if (expandedNodeId === nodeId) setExpandedNodeId(null)
    },
    [treeStructure, updateNode, deleteNode, expandedNodeId]
  )

  // ── Duplicate node ──
  const handleDuplicate = useCallback(
    (nodeId: string) => {
      duplicateNode(nodeId)
    },
    [duplicateNode]
  )

  // ── Subtree collapse toggle ──
  const handleToggleSubtreeCollapse = useCallback((nodeId: string) => {
    setCollapsedNodeIds((prev) => {
      const next = new Set(prev)
      if (next.has(nodeId)) next.delete(nodeId)
      else next.add(nodeId)
      return next
    })
  }, [])

  // ── Convert answer stub to a real node type ──
  const handleSelectAnswerType = useCallback(
    (nodeId: string, type: 'decision' | 'action' | 'solution') => {
      updateNode(nodeId, { type })
      setExpandedNodeId(nodeId)
      selectNode(nodeId)
    },
    [updateNode, selectNode]
  )

  // ── Add node flow ──
  const handleAddNodeSelect = useCallback(
    (type: NodeType, parentId: string, optionId?: string) => {
      const newId = addNode(parentId, type)
      setNewNodeIds((prev) => new Set([...prev, newId]))
      setPendingLinks((prev) => {
        const next = new Map(prev)
        next.set(newId, { parentId, optionId })
        return next
      })
      setExpandedNodeId(newId)
      setPendingAddKey(null)
    },
    [addNode]
  )

  // ── Drag & drop ──
  const handleDragStart = useCallback(
    (e: React.DragEvent, nodeId: string) => {
      e.dataTransfer.effectAllowed = 'move'
      // Find parent and index for this node
      const findParentAndIndex = (
        searchNode: TreeStructure,
        targetId: string,
        // eslint-disable-next-line @typescript-eslint/no-unused-vars
        _parentId: string | null
      ): { parentId: string | null; index: number } | null => {
        if (searchNode.children) {
          for (let i = 0; i < searchNode.children.length; i++) {
            if (searchNode.children[i].id === targetId) {
              return { parentId: searchNode.id, index: i }
            }
            const found = findParentAndIndex(
              searchNode.children[i],
              targetId,
              searchNode.id
            )
            if (found) return found
          }
        }
        return null
      }

      if (!treeStructure) return
      const location = findParentAndIndex(treeStructure, nodeId, null)
      if (location) {
        setDragState({
          nodeId,
          parentId: location.parentId,
          index: location.index,
        })
      }
    },
    [treeStructure]
  )

  const handleDragOver = useCallback(
    (e: React.DragEvent, parentId: string | null, index: number) => {
      e.preventDefault()
      setDragOverTarget({ parentId, index })
    },
    []
  )

  const handleDrop = useCallback(
    (e: React.DragEvent, targetParentId: string | null, targetIndex: number) => {
      e.preventDefault()
      if (!dragState || !targetParentId) {
        setDragState(null)
        setDragOverTarget(null)
        return
      }

      const { parentId: sourceParentId, index: sourceIndex } = dragState

      if (sourceParentId === targetParentId) {
        const adjustedIndex =
          sourceIndex < targetIndex ? targetIndex - 1 : targetIndex
        if (sourceIndex !== adjustedIndex) {
          reorderNodes(sourceParentId!, sourceIndex, adjustedIndex)
        }
      }
      // Cross-parent move intentionally not supported in canvas (complex to handle safely)

      setDragState(null)
      setDragOverTarget(null)
    },
    [dragState, reorderNodes]
  )

  const handleDragEnd = useCallback(() => {
    setDragState(null)
    setDragOverTarget(null)
  }, [])

  // ── Recursive node renderer ──
  const renderNode = useCallback(
    (
      node: TreeStructure,
      parentId: string | null,
      index: number,
      optionLabel?: string
    ): React.ReactNode => {
      const isExpanded = expandedNodeId === node.id
      const isNew = newNodeIds.has(node.id)
      const isSubtreeCollapsed = collapsedNodeIds.has(node.id)
      const nodeChildren = node.children || []

      // For decision nodes, order children by option link order
      const orderedChildren: Array<{
        child: TreeStructure
        optionLabel?: string
        optionId?: string
        childIndex: number
      }> = []

      if (node.type === 'decision' && node.options && nodeChildren.length > 0) {
        // First: children linked by options (in option order)
        const linkedChildIds = new Set<string>()
        node.options.forEach((opt) => {
          const linked = nodeChildren.find((c) => c.id === opt.next_node_id)
          if (linked) {
            orderedChildren.push({
              child: linked,
              optionLabel: opt.label || undefined,
              optionId: opt.id,
              childIndex: nodeChildren.indexOf(linked),
            })
            linkedChildIds.add(linked.id)
          }
        })
        // Then: unlinked children
        nodeChildren.forEach((child, idx) => {
          if (!linkedChildIds.has(child.id)) {
            orderedChildren.push({
              child,
              childIndex: idx,
            })
          }
        })
      } else {
        nodeChildren.forEach((child, idx) => {
          orderedChildren.push({ child, childIndex: idx })
        })
      }

      // Determine if this node has any children to render
      const hasChildren = orderedChildren.length > 0

      // Determine "add" targets for this node
      // For decision nodes: one add-button per option (not-yet-linked options)
      // For action nodes: one add-button below
      // For solution: none
      const unlinkedOptions =
        node.type === 'decision' && node.options
          ? node.options.filter(
              (opt) =>
                opt.label.trim() &&
                (!opt.next_node_id ||
                  !nodeChildren.find((c) => c.id === opt.next_node_id))
            )
          : []

      const showSingleAddButton =
        node.type === 'action' && !hasChildren

      return (
        <div
          key={node.id}
          className="flex flex-col items-center"
          ref={(el) => {
            if (el) nodeRefs.current.set(node.id, el as HTMLDivElement)
            else nodeRefs.current.delete(node.id)
          }}
        >
          {/* Drop indicator above */}
          {dragOverTarget?.parentId === parentId &&
            dragOverTarget.index === index && (
              <div className="mb-1 h-1 w-full rounded-full bg-primary" />
            )}

          {/* Option label tag (above card, shown when this is a branch from a decision) */}
          {optionLabel && (
            <div className="mb-1 rounded bg-muted px-2 py-0.5 text-[10px] text-[#848b9b] font-sans text-xs">
              {optionLabel}
            </div>
          )}

          {/* The node card — answer stubs get their own component */}
          {node.type === 'answer' ? (
            <AnswerStubCard
              node={node}
              fromOption={optionLabel}
              onSelectType={handleSelectAnswerType}
              onDelete={handleDelete}
            />
          ) : (
            <TreeCanvasNode
              node={node}
              depth={0}
              fromOption={optionLabel}
              isExpanded={isExpanded}
              isNew={isNew}
              hasChildren={nodeChildren.length > 0}
              isSubtreeCollapsed={isSubtreeCollapsed}
              onToggleExpand={() => handleToggleExpand(node.id)}
              onToggleSubtreeCollapse={() => handleToggleSubtreeCollapse(node.id)}
              onSave={handleSave}
              onCancelNew={handleCancelNew}
              onDelete={handleDelete}
              onDuplicate={handleDuplicate}
              onDragStart={handleDragStart}
              onDragOver={(e) => handleDragOver(e, parentId, index)}
              onDrop={(e) => handleDrop(e, parentId, index)}
            />
          )}

          {/* Unlinked option add buttons (decision nodes with unlinked options) */}
          {!isExpanded && unlinkedOptions.length > 0 && (
            <div className="mt-3 flex flex-col items-center gap-2">
              {unlinkedOptions.map((opt) => {
                const key = addKey(node.id, opt.id)
                return (
                  <div key={opt.id} className="flex flex-col items-center gap-1">
                    <div className="h-4 w-px bg-border" />
                    <span className="text-[10px] text-[#848b9b] font-sans text-xs">
                      {opt.label || '(unlabeled option)'}
                    </span>
                    {pendingAddKey === key ? (
                      <AddNodePicker
                        onSelect={(type) =>
                          handleAddNodeSelect(type, node.id, opt.id)
                        }
                        onCancel={() => setPendingAddKey(null)}
                      />
                    ) : (
                      <AddNodeButton
                        label="Add child"
                        onClick={() => setPendingAddKey(key)}
                      />
                    )}
                  </div>
                )
              })}
            </div>
          )}

          {/* Single add button for action nodes without children */}
          {!isExpanded && showSingleAddButton && (
            <div className="mt-3 flex flex-col items-center gap-1">
              <div className="h-4 w-px bg-border" />
              {pendingAddKey === node.id ? (
                <AddNodePicker
                  onSelect={(type) => handleAddNodeSelect(type, node.id)}
                  onCancel={() => setPendingAddKey(null)}
                />
              ) : (
                <AddNodeButton
                  label="Add next node"
                  onClick={() => setPendingAddKey(node.id)}
                />
              )}
            </div>
          )}

          {/* Collapsed subtree pill */}
          {hasChildren && !isExpanded && isSubtreeCollapsed && (
            <div className="mt-3 flex flex-col items-center">
              <div className="h-4 w-px bg-border" />
              <button
                type="button"
                onClick={() => handleToggleSubtreeCollapse(node.id)}
                className="rounded-full border border-dashed border-[#1e2130] bg-[#14161d] px-3 py-1 text-[10px] text-[#848b9b] font-sans text-xs hover:border-primary/40 hover:text-[#e2e5eb] transition-colors"
              >
                {orderedChildren.length} node{orderedChildren.length !== 1 ? 's' : ''} hidden — click to expand
              </button>
            </div>
          )}

          {/* Connector + Children */}
          {hasChildren && !isExpanded && !isSubtreeCollapsed && (
            <div className="mt-3 flex flex-col items-center w-full">
              {/* Trunk line from card down */}
              <div className="h-4 w-px bg-border" />

              {orderedChildren.length === 1 ? (
                // Single child: straight vertical
                <div className="flex flex-col items-center">
                  {renderNode(
                    orderedChildren[0].child,
                    node.id,
                    orderedChildren[0].childIndex,
                    orderedChildren[0].optionLabel
                  )}
                </div>
              ) : (
                // Multiple children: horizontal branching
                // The fork line and child lanes share the same flex container so the
                // line is sized by the actual rendered children, not a hardcoded estimate.
                <div className="flex items-start justify-center gap-6 w-full relative"
                  style={{ maxWidth: `${orderedChildren.length * 360}px` }}
                >
                  {/* Horizontal fork line — absolutely positioned, aligned to child centers.
                      Spans from center of first lane to center of last lane. */}
                  <div
                    className="absolute top-0 h-px bg-border pointer-events-none"
                    style={{
                      left: `calc(${100 / (orderedChildren.length * 2)}%)`,
                      right: `calc(${100 / (orderedChildren.length * 2)}%)`,
                    }}
                  />

                  {orderedChildren.map(({ child, optionLabel: ol, childIndex }) => (
                    <div
                      key={child.id}
                      className="flex flex-col items-center min-w-[260px]"
                    >
                      {/* Vertical stub into child lane */}
                      <div className="h-4 w-px bg-border" />
                      {renderNode(child, node.id, childIndex, ol)}
                    </div>
                  ))}
                </div>
              )}
            </div>
          )}
        </div>
      )
    },
    [
      expandedNodeId,
      newNodeIds,
      collapsedNodeIds,
      dragOverTarget,
      handleToggleExpand,
      handleToggleSubtreeCollapse,
      handleSave,
      handleCancelNew,
      handleDelete,
      handleDuplicate,
      handleSelectAnswerType,
      handleDragStart,
      handleDragOver,
      handleDrop,
      pendingAddKey,
      handleAddNodeSelect,
    ]
  )

  // ── Empty state ──
  if (!treeStructure) {
    return (
      <div className="flex h-full items-center justify-center">
        <div className="text-center">
          <div className="mb-2 text-[#848b9b] text-sm">
            No tree structure. Start by saving a tree name.
          </div>
        </div>
      </div>
    )
  }

  return (
    <div
      className="relative h-full w-full overflow-auto"
      onDragEnd={handleDragEnd}
      style={{
        // Subtle dot grid background
        backgroundImage:
          'radial-gradient(circle, var(--color-border) 1px, transparent 1px)',
        backgroundSize: '24px 24px',
      }}
    >
      <div className="flex min-h-full min-w-full items-start justify-center p-8 pb-24">
        <div className="flex flex-col items-center">
          {/* START badge above root */}
          <div className="mb-2 rounded-full border border-[#1e2130] bg-[#14161d] px-3 py-1 text-xs font-sans text-xs text-[#848b9b]">
            START
          </div>
          <div className="mb-1 h-4 w-px bg-border" />

          {renderNode(treeStructure, null, 0)}
        </div>
      </div>
    </div>
  )
}

export default TreeCanvas