import dagre from '@dagrejs/dagre'
import type { Node, Edge } from '@xyflow/react'

const NODE_WIDTH = 280
const DEFAULT_NODE_HEIGHT = 100

interface LayoutOptions {
  direction?: 'TB' | 'LR'
  rankSep?: number
  nodeSep?: number
}

export function getLayoutedElements(
  nodes: Node[],
  edges: Edge[],
  options: LayoutOptions = {}
): Node[] {
  const { direction = 'TB', rankSep = 100, nodeSep = 40 } = options

  const g = new dagre.graphlib.Graph()
  g.setDefaultEdgeLabel(() => ({}))
  g.setGraph({ rankdir: direction, ranksep: rankSep, nodesep: nodeSep })

  nodes.forEach((node) => {
    const height = node.measured?.height ?? node.data?.estimatedHeight ?? DEFAULT_NODE_HEIGHT
    g.setNode(node.id, { width: NODE_WIDTH, height })
  })

  edges.forEach((edge) => {
    g.setEdge(edge.source, edge.target)
  })

  dagre.layout(g)

  return nodes.map((node) => {
    const dagreNode = g.node(node.id)
    // dagre gives center positions — convert to top-left for React Flow
    const x = dagreNode.x - NODE_WIDTH / 2
    const y = dagreNode.y - (dagreNode.height ?? DEFAULT_NODE_HEIGHT) / 2
    return { ...node, position: { x, y } }
  })
}

export { NODE_WIDTH, DEFAULT_NODE_HEIGHT }