package ocelot.desktop.ui.layout

import ocelot.desktop.geometry.Size2D
import ocelot.desktop.ui.widget.Widget
import ocelot.desktop.util.Orientation

import scala.collection.immutable.ArraySeq

class LinearLayout(
  widget: Widget,
  orientation: Orientation.Value = Orientation.Horizontal,
  justifyContent: JustifyContent.Value = JustifyContent.Start,
  alignItems: AlignItems.Value = AlignItems.Stretch,
  gap: Float = 0f,
) extends Layout(widget) {

  override def recalculateBounds(): Unit = {
    super.recalculateBounds()

    var minHeight, maxHeight, minWidth, maxWidth: Float = 0f
    val totalGap = if (widget.children.nonEmpty) (widget.children.length - 1) * gap else 0

    orientation match {
      case Orientation.Vertical =>
        minHeight = widget.children.map(_.minimumSize.height).sum + totalGap
        maxHeight = widget.children.map(_.maximumSize.height).sum + totalGap

        minWidth = if (widget.children.isEmpty) 0 else widget.children.map(_.minimumSize.width).max
        maxWidth = if (widget.children.isEmpty) 0 else widget.children.map(_.maximumSize.width).max

      case Orientation.Horizontal =>
        minWidth = widget.children.map(_.minimumSize.width).sum + totalGap
        maxWidth = widget.children.map(_.maximumSize.width).sum + totalGap

        minHeight = if (widget.children.isEmpty) 0 else widget.children.map(_.minimumSize.height).max
        maxHeight = if (widget.children.isEmpty) 0 else widget.children.map(_.maximumSize.height).max
    }

    minimumSize = Size2D(minWidth, minHeight)
    maximumSize = Size2D(maxWidth, maxHeight)
  }

  override def relayout(): Unit = {
    val (areaWidth, areaHeight) = (widget.size.width, widget.size.height)

    val totalSpace = orientation match {
      case Orientation.Vertical => areaHeight
      case Orientation.Horizontal => areaWidth
    }

    val stretch = alignItems == AlignItems.Stretch

    var usedSpace = 0f

    for (child <- widget.children) {
      child.recalculateBounds()
      val minSize = child.minimumSize

      if (stretch) {
        child.rawSetSize(orientation match {
          case Orientation.Vertical => Size2D(areaWidth, minSize.height)
          case Orientation.Horizontal => Size2D(minSize.width, areaHeight)
        })
      } else {
        child.rawSetSize(minSize)
      }

      usedSpace += (orientation match {
        case Orientation.Vertical => minSize.height
        case Orientation.Horizontal => minSize.width
      }) + gap
    }

    if (widget.children.nonEmpty)
      usedSpace -= gap

    var freeSpace = totalSpace - usedSpace
    var newFreeSpace = freeSpace

    do {
      freeSpace = newFreeSpace
      newFreeSpace = distributeFreeSpace(freeSpace)
    } while (newFreeSpace < freeSpace)

    setPositions(totalSpace - freeSpace)

    for (child <- widget.children)
      child.relayout()
  }

  private def distributeFreeSpace(free: Float): Float = {
    var used = 0f
    var slicesLeft = widget.children.length // TODO: Stretch factor

    for (child <- widget.children) {
      val slice = (free - used) / slicesLeft

      val oldSize = child.size
      child.rawSetSize(oldSize + (orientation match {
        case Orientation.Vertical => Size2D(0, slice)
        case Orientation.Horizontal => Size2D(slice, 0)
      }))

      child.rawSetSize(child.size.clamped(child.minimumSize, child.maximumSize))

      used += (orientation match {
        case Orientation.Vertical => child.size.height - oldSize.height
        case Orientation.Horizontal => child.size.width - oldSize.width
      })

      slicesLeft -= 1
    }

    free - used
  }

  private def setPositions(size: Float): Unit = {
    var pos = (orientation, justifyContent) match {
      case (Orientation.Vertical, JustifyContent.Start) =>
        widget.position.y
      case (Orientation.Vertical, JustifyContent.End) =>
        widget.position.y + widget.size.height - size
      case (Orientation.Vertical, JustifyContent.Center) =>
        widget.position.y + widget.size.height / 2 - size / 2

      case (Orientation.Horizontal, JustifyContent.Start) =>
        widget.position.x
      case (Orientation.Horizontal, JustifyContent.End) =>
        widget.position.x + widget.size.width - size
      case (Orientation.Horizontal, JustifyContent.Center) =>
        widget.position.x + widget.size.width / 2 - size / 2
    }

    for (child <- widget.children) {
      orientation match {
        case Orientation.Vertical =>
          child.rawSetPosition(child.position.setY(pos))
          pos += child.size.height + gap

        case Orientation.Horizontal =>
          child.rawSetPosition(child.position.setX(pos))
          pos += child.size.width + gap
      }

      val align = (child match {
        case w: LinearLayout.WithOptions => w.align
        case _ => None
      }).getOrElse(alignItems)

      (orientation, align) match {
        case (Orientation.Vertical, AlignItems.Start)
            | (Orientation.Vertical, AlignItems.Stretch) =>
          child.rawSetPosition(child.position.setX(widget.position.x))
        case (Orientation.Vertical, AlignItems.End) =>
          child.rawSetPosition(child.position.setX(widget.position.x + widget.size.width - child.size.width))
        case (Orientation.Vertical, AlignItems.Center) =>
          child.rawSetPosition(child.position.setX(widget.position.x + widget.size.width / 2 - child.size.width / 2))

        case (Orientation.Horizontal, AlignItems.Start)
            | (Orientation.Horizontal, AlignItems.Stretch) =>
          child.rawSetPosition(child.position.setY(widget.position.y))
        case (Orientation.Horizontal, AlignItems.End) =>
          child.rawSetPosition(child.position.setY(widget.position.y + widget.size.height - child.size.height))
        case (Orientation.Horizontal, AlignItems.Center) =>
          child.rawSetPosition(child.position.setY(widget.position.y + widget.size.height / 2 - child.size.height / 2))
      }

      child.rawSetPosition(child.position.round)
    }
  }
}

object LinearLayout {
  class WithOptions(widget: Widget, val align: Option[AlignItems.Value]) extends Widget {
    override protected val layout: Layout = new CopyLayout(this)

    children = ArraySeq(widget)
  }

  object WithOptions {
    def apply(
      align: Option[AlignItems.Value] = None,
    )(widget: Widget) = new WithOptions(widget, align)
  }
}
