import type { Edge, Node, Position } from '@vue-flow/core'
import dagre from '@dagrejs/dagre'

const DEFAULT_NODE_WIDTH = 380
const DEFAULT_HEADER_HEIGHT = 62
const MIN_NODE_HEIGHT = 170

type LayoutDirection = 'LR' | 'TB'

type LayoutOptions = {
  direction?: LayoutDirection
  useSavedPositions?: boolean
  dimensions?: Record<string, { width: number; height: number }>
}

function estimatePortHeight(port: any, direction: 'in' | 'out') {
  let height = 54
  if (port?.type === 'Image' && port?.previewUrl) height += 108
  if (direction === 'in') height += port?.output ? 32 : 36
  return height
}

function estimateColumnHeight(ports: any[] | undefined, direction: 'in' | 'out') {
  if (!Array.isArray(ports) || ports.length === 0) return 44
  return ports.reduce((total, port) => total + estimatePortHeight(port, direction), 0) + Math.max(0, ports.length - 1) * 7
}

function estimateNodeHeight(node: Node) {
  const inputHeight = estimateColumnHeight(node.data?.inputs, 'in')
  const outputHeight = estimateColumnHeight(node.data?.outputs, 'out')
  return Math.max(MIN_NODE_HEIGHT, DEFAULT_HEADER_HEIGHT + Math.max(inputHeight, outputHeight) + 26)
}

function targetPosition(direction: LayoutDirection): Position {
  return direction === 'LR' ? 'left' as Position : 'top' as Position
}

function sourcePosition(direction: LayoutDirection): Position {
  return direction === 'LR' ? 'right' as Position : 'bottom' as Position
}

export function layoutGraph(nodes: Node[], edges: Edge[], options: LayoutOptions = {}) {
  const direction = options.direction || 'LR'

  if (options.useSavedPositions && nodes.every((node) => node.data?.hasSavedPosition)) {
    return nodes.map((node) => ({
      ...node,
      targetPosition: targetPosition(direction),
      sourcePosition: sourcePosition(direction),
    }))
  }

  const graph = new dagre.graphlib.Graph()
  graph.setDefaultEdgeLabel(() => ({}))
  graph.setGraph({
    rankdir: direction,
    nodesep: 190,
    ranksep: 260,
    marginx: 100,
    marginy: 100,
  })

  for (const node of nodes) {
    const measured = options.dimensions?.[node.id]
    graph.setNode(node.id, {
      width: Math.max(DEFAULT_NODE_WIDTH, measured?.width || DEFAULT_NODE_WIDTH),
      height: Math.max(estimateNodeHeight(node), measured?.height || 0),
    })
  }

  for (const edge of edges) {
    graph.setEdge(edge.source, edge.target)
  }

  dagre.layout(graph)

  return nodes.map((node) => {
    const graphNode = graph.node(node.id)
    const measured = options.dimensions?.[node.id]
    const width = Math.max(DEFAULT_NODE_WIDTH, measured?.width || DEFAULT_NODE_WIDTH)
    const height = Math.max(estimateNodeHeight(node), measured?.height || 0)

    return {
      ...node,
      targetPosition: targetPosition(direction),
      sourcePosition: sourcePosition(direction),
      position: graphNode
        ? { x: graphNode.x - width / 2, y: graphNode.y - height / 2 }
        : node.position,
    }
  })
}