import { useMemo, useCallback, useState, useRef, useEffect } from 'react'
import type { Node, Edge } from '@xyflow/react'
import type { TreeStructure } from '@/types'
import { getLayoutedElements, NODE_WIDTH } from '@/lib/dagreLayout'
import type { FlowCanvasNodeData } from './FlowCanvasNode'
import type { FlowCanvasAnswerNodeData } from './FlowCanvasAnswerNode'
import { useTreeEditorStore } from '@/store/treeEditorStore'

const MAX_EDGE_LABEL_LENGTH = 35

function truncateLabel(label: string): string {
  if (label.length <= MAX_EDGE_LABEL_LENGTH) return label
  return label.slice(0, MAX_EDGE_LABEL_LENGTH).trimEnd() + '…'
}

function estimateNodeHeight(node: TreeStructure): number {
  let height = 52 // header baseline
  if (node.type === 'decision' && node.options) {
    height += 24 // options header line
    height += Math.min(node.options.length, 3) * 18 // option rows (max 3 shown)
    if (node.options.length > 3) height += 18 // "+N more" row
  }
  if ((node.type === 'action' || node.type === 'solution') && node.description) {
    height += 36 // description preview (2 lines)
  }
  if (node.type === 'answer') {
    return 70 // fixed height for answer stubs
  }
  return height
}

/** Collect all node IDs in the tree. */
function collectAllIds(root: TreeStructure): Set<string> {
  const ids = new Set<string>()
  function walk(node: TreeStructure) {
    ids.add(node.id)
    node.children?.forEach(walk)
  }
  walk(root)
  return ids
}

/** Find all cross-reference edges (next_node_id pointing outside children). */
function collectCrossRefEdges(root: TreeStructure): Array<{ source: string; target: string; label?: string }> {
  const refs: Array<{ source: string; target: string; label?: string }> = []
  const allIds = collectAllIds(root)

  function walk(node: TreeStructure) {
    const childIds = new Set(node.children?.map(c => c.id) ?? [])

    // Decision options pointing outside children
    if (node.type === 'decision' && node.options) {
      for (const opt of node.options) {
        if (opt.next_node_id && !childIds.has(opt.next_node_id) && allIds.has(opt.next_node_id)) {
          refs.push({ source: node.id, target: opt.next_node_id, label: opt.label })
        }
      }
    }

    // Action next_node_id pointing to non-child (always a cross-ref since actions use next_node_id not children)
    if (node.type === 'action' && node.next_node_id && allIds.has(node.next_node_id) && !childIds.has(node.next_node_id)) {
      refs.push({ source: node.id, target: node.next_node_id, label: 'loops back' })
    }

    node.children?.forEach(walk)
  }

  walk(root)
  return refs
}

interface UseTreeLayoutResult {
  nodes: Node[]
  edges: Edge[]
  collapsedNodeIds: Set<string>
  toggleCollapse: (nodeId: string) => void
  onNodesMeasured: (measuredNodes: Node[]) => void
}

export function useTreeLayout(): UseTreeLayoutResult {
  const treeStructure = useTreeEditorStore(s => s.treeStructure)
  const validationErrors = useTreeEditorStore(s => s.validationErrors)
  const [collapsedNodeIds, setCollapsedNodeIds] = useState<Set<string>>(new Set())
  const [measuredHeights, setMeasuredHeights] = useState<Map<string, number>>(new Map())
  const correctionDone = useRef(false)

  const toggleCollapse = useCallback((nodeId: string) => {
    setCollapsedNodeIds(prev => {
      const next = new Set(prev)
      if (next.has(nodeId)) next.delete(nodeId)
      else next.add(nodeId)
      return next
    })
  }, [])

  // Convert tree structure to flat nodes and edges
  const { rawNodes, rawEdges } = useMemo(() => {
    const nodes: Node[] = []
    const edges: Edge[] = []

    if (!treeStructure) return { rawNodes: nodes, rawEdges: edges }

    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    function walk(node: TreeStructure, _parentId?: string | null) {
      const isCollapsed = collapsedNodeIds.has(node.id)
      const hasChildren = (node.children?.length ?? 0) > 0
      const hasErrors = validationErrors.some(e => e.nodeId === node.id && e.severity === 'error')

      const estimatedHeight = measuredHeights.get(node.id) ?? estimateNodeHeight(node)

      if (node.type === 'answer') {
        nodes.push({
          id: node.id,
          type: 'answerStub',
          position: { x: 0, y: 0 }, // dagre will set this
          data: {
            node,
            onSelectType: () => {}, // placeholder — set by FlowCanvas
          } satisfies FlowCanvasAnswerNodeData,
          style: { width: NODE_WIDTH },
          measured: { width: NODE_WIDTH, height: estimatedHeight },
        })
      } else {
        nodes.push({
          id: node.id,
          type: 'flowNode',
          position: { x: 0, y: 0 },
          data: {
            node,
            hasChildren,
            isCollapsed,
            hasValidationErrors: hasErrors,
            isNew: false,
            onToggleCollapse: () => {}, // placeholder — set by FlowCanvas
          } satisfies FlowCanvasNodeData,
          style: { width: NODE_WIDTH },
          measured: { width: NODE_WIDTH, height: estimatedHeight },
        })
      }

      // Skip children if collapsed
      if (isCollapsed) return

      // Create edges and recurse into children
      if (node.children) {
        // For decision nodes: order children by option link, then unlinked
        const orderedChildren = orderChildren(node)

        for (const { child, optionLabel } of orderedChildren) {
          const edgeLabel = optionLabel ? truncateLabel(optionLabel) : undefined

          edges.push({
            id: `${node.id}->${child.id}`,
            source: node.id,
            target: child.id,
            type: 'smoothstep',
            label: edgeLabel,
            labelStyle: { fill: 'hsl(var(--muted-foreground))', fontSize: 11 },
            labelBgStyle: { fill: 'hsl(var(--card))', fillOpacity: 0.9 },
            labelBgPadding: [4, 2] as [number, number],
            style: { stroke: 'hsl(var(--border))' },
          })

          walk(child, node.id)
        }
      }
    }

    walk(treeStructure, null)

    // Add cross-reference edges (dashed, purple)
    if (treeStructure) {
      const crossRefs = collectCrossRefEdges(treeStructure)
      for (const ref of crossRefs) {
        // Only add if both source and target nodes are visible (not collapsed away)
        const sourceVisible = nodes.some(n => n.id === ref.source)
        const targetVisible = nodes.some(n => n.id === ref.target)
        if (sourceVisible && targetVisible) {
          edges.push({
            id: `xref-${ref.source}->${ref.target}`,
            source: ref.source,
            target: ref.target,
            type: 'smoothstep',
            animated: true,
            label: ref.label ? truncateLabel(ref.label) : undefined,
            labelStyle: { fill: 'hsl(var(--primary))', fontSize: 10, fontWeight: 500 },
            labelBgStyle: { fill: 'hsl(var(--card))', fillOpacity: 0.95 },
            labelBgPadding: [4, 2] as [number, number],
            style: {
              stroke: 'hsl(var(--primary))',
              strokeWidth: 2,
              strokeDasharray: '6 3',
            },
            markerEnd: {
              type: 'arrowclosed' as const,
              color: 'hsl(var(--primary))',
              width: 16,
              height: 16,
            },
          })
        }
      }
    }

    return { rawNodes: nodes, rawEdges: edges }
  }, [treeStructure, collapsedNodeIds, validationErrors, measuredHeights])

  // Run dagre layout
  const { nodes, edges } = useMemo(() => {
    if (rawNodes.length === 0) return { nodes: rawNodes, edges: rawEdges }
    const layouted = getLayoutedElements(rawNodes, rawEdges)
    return { nodes: layouted, edges: rawEdges }
  }, [rawNodes, rawEdges])

  // Height measurement correction callback
  const onNodesMeasured = useCallback((measuredNodes: Node[]) => {
    if (correctionDone.current) return

    let needsCorrection = false
    const newHeights = new Map(measuredHeights)

    for (const mNode of measuredNodes) {
      const actual = mNode.measured?.height
      if (!actual) continue
      const estimated = measuredHeights.get(mNode.id) ?? estimateNodeHeight(
        (mNode.data as unknown as FlowCanvasNodeData)?.node ?? (mNode.data as unknown as FlowCanvasAnswerNodeData)?.node
      )
      if (Math.abs(actual - estimated) > 10) {
        newHeights.set(mNode.id, actual)
        needsCorrection = true
      }
    }

    if (needsCorrection) {
      correctionDone.current = true
      setMeasuredHeights(newHeights)
    }
  }, [measuredHeights])

  // Reset correction flag when tree structure changes
  useEffect(() => {
    correctionDone.current = false
  }, [treeStructure, collapsedNodeIds])

  return { nodes, edges, collapsedNodeIds, toggleCollapse, onNodesMeasured }
}

// Helper: order children by decision option links
function orderChildren(node: TreeStructure): Array<{ child: TreeStructure; optionLabel?: string }> {
  if (!node.children || node.children.length === 0) return []

  if (node.type === 'decision' && node.options) {
    const linked: Array<{ child: TreeStructure; optionLabel: string }> = []
    const linkedIds = new Set<string>()

    for (const opt of node.options) {
      if (opt.next_node_id) {
        const child = node.children.find(c => c.id === opt.next_node_id)
        if (child) {
          linked.push({ child, optionLabel: opt.label })
          linkedIds.add(child.id)
        }
      }
    }

    const unlinked = node.children
      .filter(c => !linkedIds.has(c.id))
      .map(child => ({ child, optionLabel: undefined }))

    return [...linked, ...unlinked]
  }

  return node.children.map(child => ({ child, optionLabel: undefined }))
}
