local Constraints = require "lib.simple_ui.constraints"

--- @class Flex : MultiChildElement
--- @field direction "horizontal" | "vertical"
--- @field mainAxisSize "max" | "min"
--- @field mainAxisAlignment "start" | "center" | "end"
local flex = setmetatable({}, require "lib.simple_ui.multi_child_element")
flex.__index = flex
flex.type = "Flex"
flex.direction = "horizontal"
flex.mainAxisSize = "max"
flex.mainAxisAlignment = "start"

function flex: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

return flex