#include "ui_defs.lua"

function Ui.__GridCreateRow(parent, items, idx)
    local rowH = tableMax(items, function(i1, i2) return i1.h > i2.h end).h
    local rowW = tableAccumulate(items, function (s, it) return s + it.h + parent.columsIndent end, -parent.columsIndent)

    local row = 
    {
        parent = parent,
        parentIdx = idx or 0,

        w = rowW,
        h = rowH,

        sourceItems = items,

        count = function(self)
            return #self.sourceItems
        end,

        selectedIdx = 0,

        selectItem = function(self, idx)
            if self.selectedIdx == idx then
                return
            end

            self.selectedIdx = idx
            self:onSelectionChanged(self.sourceItems[idx])
        end,

        selectNext = function(self)
            if self.selectedIdx ~= #self.sourceItems then
                local idx = clamp(self.selectedIdx + 1, 1, #self.sourceItems)
                self:selectItem(idx)
            end
        end,

        selectPrev = function(self)
            if self.selectedIdx ~= 1 then
                UiNavSkipUpdate()
                local idx = clamp(self.selectedIdx - 1, 1, #self.sourceItems)
                self:selectItem(idx)
            end
        end,

        onSelectionChanged = function(self, item)
            self.parent:onSelectionChanged(item)
            self.parent.selectedItem = item
        end,

        draw = function(self)
            UiPush()
                local rowH = 0
                UiBeginFrame()
                for i=1,#self.sourceItems do
                    UiBeginFrame()
                    UiPush()
                        self.sourceItems[i]:draw()
                    UiPop()

                    if UiIsMouseInRect(self.sourceItems[i].w, self.sourceItems[i].h)
                        and InputPressed("lmb")
                        and self.selectedIdx ~= i
                    then
                        if self.parent.selectedRowIdx ~= 0 then
                            self.parent.rows[self.parent.selectedRowIdx].selectedIdx = 0
                        end

                        self:onSelectionChanged(self.sourceItems[i])
                        self.selectedIdx = i
                        self.parent.selectedRowIdx = self.parentIdx
                    end

                    local fw, fh = UiEndFrame()
                    local sx, sy = UiGetScale()
                    local iw, ih = fw / sx, fh / sy

                    UiTranslate(self.parent.columsIndent + iw, 0)

                    if ih > rowH then
                        rowH = ih
                    end
                end

                local sx = UiGetScale()
                local rowW = UiEndFrame()
                rowW = rowW / sx

                self.w = rowW
                self.h = rowH
            UiPop()
        end
    }

    return row
end


function Ui.Grid(obj)
    local grid = Ui.Object:new
    {       
            columnCount = 1,

            rows = {},
            isFocused = false,
            ignoreNavigation = true,

            rowsIndent = 0,
            columsIndent = 0,

            state = UiCreateScrollBar(
            {
                w = 10,
                h = 120,
                color = {1, 1, 1, 1},

                draw = function(self)
                    UiPush()
                        UiTranslate(2, 2)
                        UiColor(unpack(self.color))
                        UiImageBox("common/rcorners_box_3.png", self.w, self.h, 3, 3)
                    UiPop()
                end,
            }, 
            {
                w = 14,
                h = 240,
                color = {0, 0, 0, 1},

                draw = function(self)
                    UiPush()
                            UiColor(unpack(self.color))
                        UiImageBox("common/rcorners_box_4.png", self.w, self.h, 4, 4)
                    UiPop()
                end,
            }),

            selectedRowIdx = 0,

            verticalInputBeh = UiCreateListInputBehaviour(),
            horizontalInputBeh = UiCreateListInputBehaviour(true),

            onSelectionChanged = function(self, item)
            end,

            append = function(self, item)
                Ui.Object.append(self, item)

                if #self.rows == 0 then
                    local newRow = Ui.__GridCreateRow(self, { item }, 1)
                    table.insert(self.rows, newRow)
                    return
                end

                local lastRow = self.rows[#self.rows]
                if lastRow:count() < self.columnCount then
                    table.insert(lastRow.sourceItems, item)
                else
                    local newRow = Ui.__GridCreateRow(self, { item }, #self.rows + 1)
                    table.insert(self.rows, newRow)
                end
            end,

            appendMany = function(self, items)
                for i=1,#items do
                    self:append(items[i])
                end
            end,

            selectItem = function(self, item)
                local fItem = nil
                local fItemIdx = 0
                local fRow, fRowIdx = tableFirst(
                    self.rows, 
                    function(row)
                        fItem, fItemIdx = tableFirst(row.sourceItems, function(si) return si == item end)
                        return fItem ~= nil
                    end)

                if fItem == nil then
                    return
                end

                fRow:selectItem(fItemIdx)

                if self.selectedRowIdx ~= 0 then
                    self.rows[self.selectedRowIdx].selectedIdx = 0
                end

                self.selectedRowIdx = fRowIdx
                self.state.navIdx = fRowIdx
            end,

            clear = function(self)
                self.children = {}
                self.rows = {}
                self.selectedItem = nil

                self:reset()
            end,

            reset = function(self)
                self.state:reset()
                self.selectedRowIdx = 0
            end,

            draw = function(self)
                UiFrameOccupy(self.w, self.h)

                self.state.contentSpan = tableAccumulate(self.rows, function(s, it) return s + it.h + self.rowsIndent end, -self.rowsIndent)
                self.state.listSpan = self.h

                self.state.rect.w = self.w
                self.state.rect.h = self.h
                self.state.rect.x, self.state.rect.y = UiGetCursorPos()
                

                if InputPressed("lmb") and UiIsMouseInRect(self.w, self.h) and self.state.allowListDrag then
                    self.state.listDrag = true
                    local _, my = UiGetMousePos()
                    self.state.listDragPos = my
                elseif self.state.listDrag and InputDown("lmb") then
                    local _, my = UiGetMousePos()
                    self.state:addOffset(my - self.state.listDragPos, self.h, self.contentHeight, true)
                    self.state.listDragPos = my
                elseif InputReleased("lmb") then
                    self.state.listDrag = false
                end

                self.verticalInputBeh:update()
                self.horizontalInputBeh:update()

                -- if the list wasn't drawn on previous draw then return list to initial state
                if self.state.lastFrame ~= UiGetFrameNo() - 1 and not self.state.preserveState then
                    self.state:reset()
                end
                self.state.lastFrame = UiGetFrameNo()

                UiPush()
                    UiTranslate(self.state:getTranslateToScrollbar())
                    self.state:drawScrollbar()
                UiPop()
                
                UiPush()
                    UiFrameSkipItem(true)
                    UiAlign("top left")
                    UiClipRect(self.w, self.h, true)

                    UiPush()
                        if UiIsMouseInRect(self.w, self.h) and not self.state.ignoreWheelScroll then
                            local mouseScroll = InputValue("mousewheel") * self.state.scrollSpeed
                            if mouseScroll ~= 0 then
                                self.state:addOffset(mouseScroll)
                                self.state.navIdx = 0
                            end
                        end

                        UiTranslate(self.state:getTranslateOffset())

                        local focusedComponentIdx = 0
                        for i=1,#self.rows do
                            local row = self.rows[i]
                            
                            local sx, sy = UiGetScale()
                            local isVisible = not UiIsFullyClipped(self.rows[i].w * sx, self.rows[i].h * sy)

                            if isVisible then
                                UiPush()
                                    row:draw()
                                UiPop()
                            end

                            UiTranslate(0, row.h + self.rowsIndent)
                        end

                        self:__resolveSelection()

                        if self.state.navIdx ~= 0 then
                            local listDim = self.h
                            local itemPos = self.state:getOffset()

                            for i=1,self.state.navIdx - 1 do
                                itemPos = itemPos + self.rows[i].h + self.rowsIndent
                            end

                            local topBorder = (self.state.navIdx == 1) and 0 or (0.1 * listDim * UiGetScale())
                            local bottomBorder = (self.state.navIdx == #self.rows) and (listDim) or (listDim * 0.9 * UiGetScale())
                            local itemInVisibleRegion = (itemPos >= topBorder) and ((itemPos + self.rows[self.state.navIdx].h * UiGetScale()) <= bottomBorder)

                            if not itemInVisibleRegion and not self.state.lerping then
                                if itemPos < topBorder then
                                    local offset = topBorder - itemPos
                                    self.state:addOffset(offset / UiGetScale(), not self.state.allowLerp)
                                    self.state.allowLerp = true
                                elseif itemPos + self.rows[self.state.navIdx].h > bottomBorder then
                                    local offset = bottomBorder - (itemPos + self.rows[self.state.navIdx].h)
                                    self.state:addOffset(offset / UiGetScale(), not self.state.allowLerp)
                                    self.state.allowLerp = true
                                end
                            end
                        end
                    UiPop()
                    UiNavGroupEnd()
                UiPop()

                self.state.selectedIdx = focusedComponentIdx
            end,

            __resolveSelection = function(self)
                if not self.rows[1] or LastInputDevice() ~= UI_DEVICE_GAMEPAD or not self.isFocused then
                    return
                end

                if self.selectedRowIdx == 0 then
                    self.selectedRowIdx = 1
                    self.rows[1]:selectItem(1)
                    self.state.navIdx = 1
                end

                if self.verticalInputBeh.nextRequested and not self.state.lerping then
                    if self.selectedRowIdx ~= #self.rows then
                        local ri = clamp(self.selectedRowIdx + 1, 1, #self.rows)
                        local col = self.rows[self.selectedRowIdx].selectedIdx
                        if self.rows[ri].sourceItems[col] ~= nil then
                            self.rows[ri]:selectItem(col)
                            self.rows[self.selectedRowIdx].selectedIdx = 0
                            self.selectedRowIdx = ri

                            self.state.navIdx = ri
                        end
                    end
                elseif self.verticalInputBeh.prevRequested and not self.state.lerping then
                    if self.selectedRowIdx ~= 1 then
                        local ri = clamp(self.selectedRowIdx - 1, 1, #self.rows)
                        local col = self.rows[self.selectedRowIdx].selectedIdx
                        if self.rows[ri].sourceItems[col] ~= nil then
                            self.rows[ri]:selectItem(col)
                            self.rows[self.selectedRowIdx].selectedIdx = 0
                            self.selectedRowIdx = ri

                            self.state.navIdx = ri
                        end
                    end
                elseif self.horizontalInputBeh.nextRequested and not self.state.lerping then
                    local row = self.rows[self.selectedRowIdx]
                    row:selectNext()
                elseif self.horizontalInputBeh.prevRequested and not self.state.lerping then
                    local row = self.rows[self.selectedRowIdx]
                    row:selectPrev()
                end
            end
        }

    return grid:new(obj)
end