dfhack/plugins/lua/buildingplan/itemselection.lua

local _ENV = mkmodule('plugins.buildingplan.itemselection')

local gui = require('gui')
local pens = require('plugins.buildingplan.pens')
local utils = require('utils')
local widgets = require('gui.widgets')

local uibs = df.global.buildreq
local to_pen = dfhack.pen.parse

local BUILD_TEXT_PEN = to_pen{fg=COLOR_BLACK, bg=COLOR_GREEN, keep_lower=true}
local BUILD_TEXT_HPEN = to_pen{fg=COLOR_WHITE, bg=COLOR_GREEN, keep_lower=true}

-- map of building type -> {set=set of recently used, list=list of recently used}
-- most recent entries are at the *end* of the list
local recently_used = {}

local function sort_by_type(a, b)
    local ad, bd = a.data, b.data
    return ad.item_type < bd.item_type or
            (ad.item_type == bd.item_type and ad.item_subtype < bd.item_subtype) or
            (ad.item_type == bd.item_type and ad.item_subtype == bd.item_subtype and a.search_key < b.search_key) or
            (ad.item_type == bd.item_type and ad.item_subtype == bd.item_subtype and a.search_key == b.search_key and ad.quality > bd.quality)
end

local function sort_by_recency(a, b)
    local tracker = recently_used[uibs.building_type]
    if not tracker then return sort_by_type(a, b) end
    local recent_a, recent_b = tracker.set[a.search_key], tracker.set[b.search_key]
    -- if they're both in the set, return the one with the greater index,
    -- indicating more recent
    if recent_a and recent_b then return recent_a > recent_b end
    if recent_a and not recent_b then return true end
    if not recent_a and recent_b then return false end
    return sort_by_type(a, b)
end

local function sort_by_name(a, b)
    return a.search_key < b.search_key or
            (a.search_key == b.search_key and sort_by_type(a, b))
end

local function sort_by_quantity(a, b)
    local ad, bd = a.data, b.data
    return ad.quantity > bd.quantity or
            (ad.quantity == bd.quantity and sort_by_type(a, b))
end

ItemSelection = defclass(ItemSelection, widgets.Window)
ItemSelection.ATTRS{
    frame_title='Choose items',
    frame={w=56, h=20, l=4, t=8},
    resizable=true,
    index=DEFAULT_NIL,
    desc=DEFAULT_NIL,
    quantity=DEFAULT_NIL,
    on_submit=DEFAULT_NIL,
    on_cancel=DEFAULT_NIL,
}

function ItemSelection:init()
    self.num_selected = 0
    self.selected_set = {}
    local plural = self.quantity == 1 and '' or 's'

    self:addviews{
        widgets.Label{
            frame={t=0, l=0, r=10},
            text={
                self.desc,
                plural,
                NEWLINE,
                ('Select up to %d item%s ('):format(self.quantity, plural),
                {text=function() return self.num_selected end},
                ' selected)',
            },
        },
        widgets.Label{
            frame={r=0, w=11, t=0, h=3},
            text_pen=BUILD_TEXT_PEN,
            text_hpen=BUILD_TEXT_HPEN,
            text={
                '           ', NEWLINE,
                '  Confirm  ', NEWLINE,
                '           ',
            },
            on_click=self:callback('submit'),
        },
        widgets.FilteredList{
            view_id='flist',
            frame={t=3, l=0, r=0, b=4},
            case_sensitive=false,
            choices=self:get_choices(sort_by_recency),
            icon_width=2,
            on_submit=self:callback('toggle_group'),
            edit_on_char=function(ch) return ch:match('[%l -]') end,
        },
        widgets.CycleHotkeyLabel{
            frame={l=0, b=2},
            key='CUSTOM_SHIFT_R',
            label='Sort by:',
            options={
                {label='Recently used', value=sort_by_recency},
                {label='Name', value=sort_by_name},
                {label='Amount', value=sort_by_quantity},
            },
            on_change=self:callback('on_sort'),
        },
        widgets.HotkeyLabel{
            frame={l=0, b=1},
            key='SELECT',
            label='Use all/none',
            auto_width=true,
            on_activate=function() self:toggle_group(self.subviews.flist.list:getSelected()) end,
        },
        widgets.HotkeyLabel{
            frame={l=22, b=1},
            key='CUSTOM_SHIFT_C',
            label='Confirm',
            auto_width=true,
            on_activate=self:callback('submit'),
        },
        widgets.HotkeyLabel{
            frame={l=38, b=1},
            key='LEAVESCREEN',
            label='Go back',
            auto_width=true,
            on_activate=self:callback('on_cancel'),
        },
        widgets.HotkeyLabel{
            frame={l=0, b=0},
            key='KEYBOARD_CURSOR_RIGHT_FAST',
            key_sep='    : ',
            label='Use one',
            auto_width=true,
            on_activate=function() self:increment_group(self.subviews.flist.list:getSelected()) end,
        },
        widgets.Label{
            frame={l=6, b=0, w=5},
            text_pen=COLOR_LIGHTGREEN,
            text='Right',
        },
        widgets.HotkeyLabel{
            frame={l=23, b=0},
            key='KEYBOARD_CURSOR_LEFT_FAST',
            key_sep='   : ',
            label='Use one fewer',
            auto_width=true,
            on_activate=function() self:decrement_group(self.subviews.flist.list:getSelected()) end,
        },
        widgets.Label{
            frame={l=29, b=0, w=4},
            text_pen=COLOR_LIGHTGREEN,
            text='Left',
        },
    }
end

-- resort and restore selection
function ItemSelection:on_sort(sort_fn)
    local flist = self.subviews.flist
    local saved_filter = flist:getFilter()
    flist:setFilter('')
    flist:setChoices(self:get_choices(sort_fn), flist:getSelected())
    flist:setFilter(saved_filter)
end

local function make_search_key(str)
    local out = ''
    for c in str:gmatch("[%w%s]") do
        out = out .. c
    end
    return out
end

function ItemSelection:get_choices(sort_fn)
    local item_ids = require('plugins.buildingplan').getAvailableItems(uibs.building_type,
            uibs.building_subtype, uibs.custom_type, self.index-1)
    local buckets = {}
    for _,item_id in ipairs(item_ids) do
        local item = df.item.find(item_id)
        if not item then goto continue end
        local desc = dfhack.items.getDescription(item, 0, true)
        if buckets[desc] then
            local bucket = buckets[desc]
            table.insert(bucket.data.item_ids, item_id)
            bucket.data.quantity = bucket.data.quantity + 1
        else
            local entry = {
                search_key=make_search_key(desc),
                icon=self:callback('get_entry_icon', item_id),
                data={
                    item_ids={item_id},
                    item_type=item:getType(),
                    item_subtype=item:getSubtype(),
                    quantity=1,
                    quality=item:getQuality(),
                    selected=0,
                },
            }
            buckets[desc] = entry
        end
        ::continue::
    end
    local choices = {}
    for desc,choice in pairs(buckets) do
        local data = choice.data
        choice.text = {
            {width=10, text=function() return ('[%d/%d]'):format(data.selected, data.quantity) end},
            {gap=2, text=desc},
        }
        table.insert(choices, choice)
    end
    table.sort(choices, sort_fn)
    return choices
end

function ItemSelection:increment_group(idx, choice)
    local data = choice.data
    if self.quantity <= self.num_selected then return false end
    if data.selected >= data.quantity then return false end
    data.selected = data.selected + 1
    self.num_selected = self.num_selected + 1
    local item_id = data.item_ids[data.selected]
    self.selected_set[item_id] = true
    return true
end

function ItemSelection:decrement_group(idx, choice)
    local data = choice.data
    if data.selected <= 0 then return false end
    local item_id = data.item_ids[data.selected]
    self.selected_set[item_id] = nil
    self.num_selected = self.num_selected - 1
    data.selected = data.selected - 1
    return true
end

function ItemSelection:toggle_group(idx, choice)
    local data = choice.data
    if data.selected > 0 then
        while self:decrement_group(idx, choice) do end
    else
        while self:increment_group(idx, choice) do end
    end
end

function ItemSelection:get_entry_icon(item_id)
    return self.selected_set[item_id] and pens.SELECTED_ITEM_PEN or nil
end

local function track_recently_used(choices)
    -- use same set for all subtypes
    local tracker = ensure_key(recently_used, uibs.building_type)
    for _,choice in ipairs(choices) do
        local data = choice.data
        if data.selected <= 0 then goto continue end
        local key = choice.search_key
        local recent_set = ensure_key(tracker, 'set')
        local recent_list = ensure_key(tracker, 'list')
        if recent_set[key] then
            if recent_list[#recent_list] ~= key then
                for i,v in ipairs(recent_list) do
                    if v == key then
                        table.remove(recent_list, i)
                        table.insert(recent_list, key)
                        break
                    end
                end
                tracker.set = utils.invert(recent_list)
            end
        else
            -- only keep most recent 10
            if #recent_list >= 10 then
                -- remove least recently used from list and set
                recent_set[table.remove(recent_list, 1)] = nil
            end
            table.insert(recent_list, key)
            recent_set[key] = #recent_list
        end
        ::continue::
    end
end

function ItemSelection:submit()
    local selected_items = {}
    for item_id in pairs(self.selected_set) do
        table.insert(selected_items, item_id)
    end
    if #selected_items > 0 then
        track_recently_used(self.subviews.flist:getChoices())
    end
    self.on_submit(selected_items)
end

function ItemSelection:onInput(keys)
    if keys.LEAVESCREEN or keys._MOUSE_R_DOWN then
        self.on_cancel()
        return true
    elseif keys._MOUSE_L_DOWN then
        local list = self.subviews.flist.list
        local idx = list:getIdxUnderMouse()
        if idx then
            list:setSelected(idx)
            local modstate = dfhack.internal.getModstate()
            if modstate & 2 > 0 then -- ctrl
                local choice = list:getChoices()[idx]
                if modstate & 1 > 0 then -- shift
                    self:decrement_group(idx, choice)
                else
                    self:increment_group(idx, choice)
                end
                return true
            end
        end
    end
    return ItemSelection.super.onInput(self, keys)
end

ItemSelectionScreen = defclass(ItemSelectionScreen, gui.ZScreen)
ItemSelectionScreen.ATTRS {
    focus_path='dwarfmode/Building/Placement/dfhack/lua/buildingplan/itemselection',
    force_pause=true,
    pass_movement_keys=true,
    pass_pause=false,
    pass_mouse_clicks=false,
    defocusable=false,
    index=DEFAULT_NIL,
    desc=DEFAULT_NIL,
    quantity=DEFAULT_NIL,
    on_submit=DEFAULT_NIL,
    on_cancel=DEFAULT_NIL,
}

function ItemSelectionScreen:init()
    self:addviews{
        ItemSelection{
            index=self.index,
            desc=self.desc,
            quantity=self.quantity,
            on_submit=self.on_submit,
            on_cancel=self.on_cancel,
        }
    }
end

return _ENV