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 = {}

function get_automaterial_selection(building_type)
    local tracker = recently_used[building_type]
    if not tracker or not tracker.list then return end
    return tracker.list[#tracker.list]
end

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=24, l=4, t=7},
    resizable=true,
    index=DEFAULT_NIL,
    desc=DEFAULT_NIL,
    quantity=DEFAULT_NIL,
    autoselect=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'

    local choices = self:get_choices(sort_by_recency)

    if self.autoselect then
        self:do_autoselect(choices)
        if self.num_selected >= self.quantity then
            self:submit(choices)
            return
        end
    end

    self:addviews{
        widgets.Panel{
            view_id='header',
            frame={t=0, h=3},
            subviews={
                widgets.Label{
                    frame={t=0, l=0, r=16},
                    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=15, t=0, h=3},
                    text_pen=BUILD_TEXT_PEN,
                    text_hpen=BUILD_TEXT_HPEN,
                    text={
                        '   Use filter  ', NEWLINE,
                        ' for remaining ', NEWLINE,
                        '     items     ',
                    },
                    on_click=self:callback('submit'),
                    visible=function() return self.num_selected < self.quantity end,
                },
                widgets.Label{
                    frame={r=0, w=15, t=0, h=3},
                    text_pen=BUILD_TEXT_PEN,
                    text_hpen=BUILD_TEXT_HPEN,
                    text={
                        '               ', NEWLINE,
                        '    Continue   ', NEWLINE,
                        '               ',
                    },
                    on_click=self:callback('submit'),
                    visible=function() return self.num_selected >= self.quantity end,
                },
            },
        },
    }

    self:addviews{
        widgets.Panel{
            view_id='body',
            frame={t=self.subviews.header.frame.h, b=4},
            subviews={
                widgets.EditField{
                    view_id='search',
                    frame={l=1, t=0},
                    label_text='Search: ',
                    on_char=function(ch) return ch:match('[%l -]') end,
                },
                widgets.CycleHotkeyLabel{
                    frame={l=1, t=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.Panel{
                    frame={l=0, t=3, r=0, b=0},
                    frame_style=gui.INTERIOR_FRAME,
                    subviews={
                        widgets.FilteredList{
                            view_id='flist',
                            frame={t=0, b=0},
                            case_sensitive=false,
                            choices=choices,
                            icon_width=2,
                            on_submit=self:callback('toggle_group'),
                        },
                    },
                },
            },
        },
        widgets.Panel{
            view_id='footer',
            frame={l=1, r=1, b=0, h=3},
            subviews={
                widgets.HotkeyLabel{
                    frame={l=0, h=1, t=0},
                    key='KEYBOARD_CURSOR_RIGHT_FAST',
                    key_sep='----: ', -- these hypens function as "padding" to be overwritten by the next Label
                    label='Use one',
                    auto_width=true,
                    on_activate=function() self:increment_group(self.subviews.flist.list:getSelected()) end,
                },
                widgets.Label{
                    frame={l=6, w=5, t=0},
                    text_pen=COLOR_LIGHTGREEN,
                    text='Right', -- this overrides the "6----" characters from the previous HotkeyLabel
                },
                widgets.HotkeyLabel{
                    frame={l=1, h=1, t=1},
                    key='KEYBOARD_CURSOR_LEFT_FAST',
                    key_sep='---: ', -- these hypens function as "padding" to be overwritten by the next Label
                    label='Use one fewer',
                    auto_width=true,
                    on_activate=function() self:decrement_group(self.subviews.flist.list:getSelected()) end,
                },
                widgets.Label{
                    frame={l=7, w=4, t=1},
                    text_pen=COLOR_LIGHTGREEN,
                    text='Left', -- this overrides the "4---" characters from the previous HotkeyLabel
                },
                widgets.HotkeyLabel{
                    frame={l=6, t=2, h=2},
                    key='SELECT',
                    label='Use all/none',
                    auto_width=true,
                    on_activate=function() self:toggle_group(self.subviews.flist.list:getSelected()) end,
                },
                widgets.HotkeyLabel{
                    frame={r=5, t=0},
                    key='LEAVESCREEN',
                    label='Go back',
                    auto_width=true,
                    on_activate=self:callback('on_cancel'),
                },
                widgets.HotkeyLabel{
                    frame={r=4, t=2},
                    key='CUSTOM_SHIFT_C',
                    label='Continue',
                    auto_width=true,
                    on_activate=self:callback('submit'),
                },
            },
        },
    }

    self.subviews.flist.list.frame.t = 0
    self.subviews.flist.edit.visible = false
    self.subviews.flist.edit = self.subviews.search
    self.subviews.search.on_change = self.subviews.flist:callback('onFilterChange')
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:do_autoselect(choices)
    if #choices == 0 then return end
    local desired = get_automaterial_selection(uibs.building_type)
    if choices[1].search_key ~= desired then return end
    self:toggle_group(1, choices[1])
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(choices)
    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(choices or self.subviews.flist:getChoices())
    end
    self.on_submit(selected_items)
end

function ItemSelection:onInput(keys)
    if keys.LEAVESCREEN or keys._MOUSE_R then
        self.on_cancel()
        return true
    elseif keys._MOUSE_L 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,
    autoselect=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,
            autoselect=self.autoselect,
            on_submit=self.on_submit,
            on_cancel=self.on_cancel,
        }
    }
end

return _ENV