import { useCallback } from 'react'
import { useDominoStoreInstance } from '../components/canvas'
import type {
  SceneElement,
  PlaceholderElement,
  Padding,
} from '../components/canvas'

export interface LayoutOptions {
  maxWidth?: number
  padding?: Padding
  colGap?: number
  rowGap?: number
}

export function useLayout(options: LayoutOptions = {}) {
  const {
    maxWidth = 7000,
    padding = { top: 100, left: 100, bottom: 100, right: 100 },
    colGap = 20,
    rowGap = 80,
  } = options
  const store = useDominoStoreInstance()

  const topPadding = padding.top ?? 0
  const leftPadding = padding.left ?? 0

  /**
   * Rearrange all existing top-level elements
   */
  const arrangeElements = useCallback(
    (overrideOptions?: LayoutOptions) => {
      const config = {
        maxWidth,
        padding,
        colGap,
        rowGap,
        ...overrideOptions,
      }

      const state = store.getState()
      const { elements, placeholders, elementOrder } = state

      const configTop = config.padding?.top ?? topPadding
      const configLeft = config.padding?.left ?? leftPadding

      // Record history
      store.getState().takeSnapshot()

      const newPositions: Record<
        string,
        { x: number; y: number; rotation: number }
      > = {}
      let currentX = configLeft
      let currentY = configTop
      let maxHeightInRow = 0

      const elementsToArrange = elementOrder
        .map((id: string) => elements[id] || placeholders[id])
        .filter(
          (el: SceneElement | undefined): el is SceneElement =>
            el !== undefined,
        )

      elementsToArrange.forEach((el: SceneElement) => {
        if (currentX + el.width > config.maxWidth) {
          currentX = configLeft
          currentY += maxHeightInRow + config.rowGap
          maxHeightInRow = 0
        }

        newPositions[el.id] = {
          x: currentX,
          y: currentY,
          rotation: 0,
        }

        maxHeightInRow = Math.max(maxHeightInRow, el.height)
        currentX += el.width + config.colGap
      })

      store.setState(
        (state: {
          elements: Record<string, SceneElement>
          placeholders: Record<string, PlaceholderElement>
        }) => {
          Object.entries(newPositions).forEach(([id, pos]) => {
            if (state.elements[id]) {
              state.elements[id] = { ...state.elements[id], ...pos }
            } else if (state.placeholders[id]) {
              state.placeholders[id] = { ...state.placeholders[id], ...pos }
            }
          })
        },
      )
    },
    [maxWidth, topPadding, leftPadding, colGap, rowGap, store, padding],
  )

  /**
   * Add a new element to the end of the flow
   */
  const addElementToFlow = useCallback(
    <T extends SceneElement>(elementTemplate: T): T => {
      const state = store.getState()
      const { elements, placeholders, elementOrder } = state

      const getAnyElement = (id: string) => elements[id] || placeholders[id]

      // Get all top-level elements to find the current "flow" tail
      const existingElements = elementOrder
        .map((id: string) => getAnyElement(id))
        .filter(
          (el: SceneElement | undefined): el is SceneElement =>
            !!el &&
            (el.type !== 'artboard' ||
              (el.type === 'artboard' && !el.parentId)),
        )

      const lastElement = existingElements[existingElements.length - 1]
      let currentX = lastElement
        ? lastElement.x + lastElement.width + colGap
        : leftPadding
      let currentY = lastElement ? lastElement.y : topPadding

      // To find the current row's max height, we look at elements on the same Y
      const currentRowElements = existingElements.filter(
        (el: SceneElement) => el.y === currentY,
      )
      const maxHeightInRow =
        currentRowElements.length > 0
          ? Math.max(...currentRowElements.map((el: SceneElement) => el.height))
          : 0

      // Wrap check
      if (currentX + elementTemplate.width > maxWidth) {
        currentX = leftPadding
        currentY += maxHeightInRow + rowGap
      }

      const newElement: SceneElement = {
        ...elementTemplate,
        x: currentX,
        y: currentY,
      }

      if (newElement.type === 'placeholder') {
        state.addPlaceholder(newElement as PlaceholderElement)
      } else {
        state.addElement(newElement)
      }

      state.setFocusedElementId(newElement.id)
      state.setSelectedIds([])

      return newElement as T
    },
    [colGap, rowGap, maxWidth, topPadding, leftPadding, store],
  )

  return { arrangeElements, addElementToFlow }
}
