local Constraints = require "lib.simple_ui.core.constraints"
local MultiChildElement = require "lib.simple_ui.core.multi_child_element"

--- @class Flex : MultiChildElement
--- @field direction Axis
--- @field mainAxisSize MainAxisSize
--- @field mainAxisAlignment MainAxisAlignment
local element = setmetatable({}, require "lib.simple_ui.core.multi_child_element")
element.__index = element
element.type = "Flex"
element.direction = "horizontal"
element.mainAxisSize = "max"
element.mainAxisAlignment = "start"

function element:layout()
    local mainAxisSize = 0
    local crossAxisSize = 0
    if self.direction == "horizontal" then
        for _, child in ipairs(self.children) do
            child._constraints_ = Constraints { maxHeight = self._constraints_.maxHeight }
            child:layout()
            if child._size_.y > crossAxisSize then crossAxisSize = child._size_.y end
            mainAxisSize = mainAxisSize + child._size_.x
        end

        local start = 0
        if self.mainAxisAlignment == "center" then
            start = self._constraints_.maxWidth / 2 - mainAxisSize / 2
        elseif self.mainAxisAlignment == "end" then
            start = self._constraints_.maxWidth - mainAxisSize
        end
        local shift = 0
        for _, child in ipairs(self.children) do
            child._offset_ = Vec3 { start + shift, 0 }
            shift = shift + child._size_.x
        end

        if self.mainAxisSize == "max" then
            self._size_ = Vec3 { self._constraints_.maxWidth, crossAxisSize }
        else
            self._size_ = Vec3 { mainAxisSize, crossAxisSize }
        end
    else
        for _, child in ipairs(self.children) do
            child._constraints_ = Constraints { maxWidth = self._constraints_.maxWidth }
            child:layout()
            if child._size_.x > crossAxisSize then crossAxisSize = child._size_.x end
            mainAxisSize = mainAxisSize + child._size_.y
        end


        local start = 0
        if self.mainAxisAlignment == "center" then
            start = self._constraints_.maxHeight / 2 - mainAxisSize / 2
        elseif self.mainAxisAlignment == "end" then
            start = self._constraints_.maxHeight - mainAxisSize
        end
        local shift = 0
        for _, child in ipairs(self.children) do
            child._offset_ = Vec3 { 0, start + shift }
            shift = shift + child._size_.y
        end

        if self.mainAxisSize == "max" then
            self._size_ = Vec3 { crossAxisSize, self._constraints_.maxHeight }
        else
            self._size_ = Vec3 { crossAxisSize, mainAxisSize }
        end
    end
end

--- @param values {direction: Axis?, mainAxisSize: MainAxisSize?, mainAxisAlignment: MainAxisAlignment?, children: UIElement[]?}
--- @return Flex
function element:new(values)
    return MultiChildElement.new(self, values)
end

return element