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 { self.offset.x + start + shift, self.offset.y }
            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 { self.offset.x, self.offset.y + 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