add support for work details and interrogation

2023-10-09 01:10:32 -07:00 · 2023-10-09 01:10:32 -07:00 · 06faeb669b
parent 060becec7c
commit 06faeb669b
1 changed files with 173 additions and 64 deletions
--- a/plugins/lua/sort/info.lua
+++ b/plugins/lua/sort/info.lua
@ -7,8 +7,9 @@ local utils = require('utils')

 local info = df.global.game.main_interface.info
 local creatures = info.creatures
-local objects = info.artifacts
 local justice = info.justice
+local objects = info.artifacts
+local work_details = info.labor.work_details

 local state = {}

@ -104,17 +105,19 @@ local function copy_to_lua_table(vec)
    return tab
 end

-local function general_search(vec, get_search_key_fn, get_sort_fn, data, filter, incremental)
+local function general_search(vec, get_search_key_fn, get_sort_fn, matches_filters_fn, data, filter, incremental)
    if not data.saved_original then
        data.saved_original = copy_to_lua_table(vec)
    elseif not incremental then
        vec:assign(data.saved_original)
    end
-    if filter ~= '' then
+    if matches_filters_fn ~= DEFAULT_NIL or filter ~= '' then
        local search_tokens = filter:split()
        for idx = #vec-1,0,-1 do
            local search_key = get_search_key_fn(vec[idx])
-            if search_key and not utils.search_text(search_key, search_tokens) then
+            if (search_key and not utils.search_text(search_key, search_tokens)) or
+                (matches_filters_fn ~= DEFAULT_NIL and not matches_filters_fn(vec[idx]))
+            then
                vec:erase(idx)
            end
        end
@ -148,30 +151,76 @@ local function make_cri_unitst_handlers(vec)
    }
 end

-local function overall_training_search(data, filter, incremental)
+local function overall_training_search(matches_filters_fn, data, filter, incremental)
    general_search(creatures.atk_index, function(elem)
            local raw = df.creature_raw.find(elem)
            if not raw then return end
            return raw.name[1]
-        end, nil, data, filter, incremental)
+        end, nil, matches_filters_fn, data, filter, incremental)
 end

-local function assign_trainer_search(data, filter, incremental)
+local function assign_trainer_search(matches_filters_fn, data, filter, incremental)
    general_search(creatures.trainer, function(elem)
            if not elem then return end
            return ('%s %s'):format(dfhack.TranslateName(elem.name), dfhack.units.getProfessionName(elem))
-        end, nil, data, filter, incremental)
+        end, nil, matches_filters_fn, data, filter, incremental)
 end

-local function interrogating_search(data, filter, incremental)
+local function get_unit_search_key(unit)
+    return ('%s %s %s'):format(
+        dfhack.units.getReadableName(unit),  -- last name is in english
+        dfhack.units.getProfessionName(unit),
+        dfhack.TranslateName(unit.name, false, true))  -- get untranslated last name
 end

-local function convicting_search(data, filter, incremental)
-    general_search(justice.conviction_list, function(elem)
-        return ('%s %s'):format(
-            dfhack.units.getReadableName(elem),
-            dfhack.units.getProfessionName(elem))
-    end, nil, data, filter, incremental)
+local function work_details_search(matches_filters_fn, data, filter, incremental)
+    if work_details.selected_work_detail_index ~= data.selected then
+        data.saved_original = nil
+        data.selected = work_details.selected_work_detail_index
+    end
+    general_search(work_details.assignable_unit, get_unit_search_key,
+        nil, matches_filters_fn, data, filter, incremental)
+end
+
+-- independent implementation of search algorithm since we need to
+-- keep two vectors in sync
+local function interrogating_search(matches_filters_fn, data, filter, incremental)
+    local vec, flags_vec = justice.interrogation_list, justice.interrogation_list_flag
+    if not data.saved_original then
+        data.saved_original = copy_to_lua_table(vec)
+        data.saved_flags = copy_to_lua_table(flags_vec)
+        data.saved_idx_map = {}
+        for idx, unit in ipairs(data.saved_original) do
+            data.saved_idx_map[unit.id] = idx  -- 1-based idx
+        end
+    else  -- sync flag changes to saved vector
+        for idx, unit in ipairs(vec) do  -- 0-based idx
+            data.saved_flags[data.saved_idx_map[unit.id]] = flags_vec[idx]
+        end
+    end
+
+    if not incremental then
+        vec:assign(data.saved_original)
+        flags_vec:assign(data.saved_flags)
+    end
+
+    if matches_filters_fn or filter ~= '' then
+        local search_tokens = filter:split()
+        for idx = #vec-1,0,-1 do
+            local search_key = get_unit_search_key(vec[idx])
+            if (search_key and not utils.search_text(search_key, search_tokens)) or
+                (matches_filters_fn and not matches_filters_fn(vec[idx], idx))
+            then
+                vec:erase(idx)
+                flags_vec:erase(idx)
+            end
+        end
+    end
+end
+
+local function convicting_search(matches_filters_fn, data, filter, incremental)
+    general_search(justice.conviction_list, get_unit_search_key,
+        nil, matches_filters_fn, data, filter, incremental)
 end

 local HANDLERS = {
@ -181,6 +230,7 @@ local HANDLERS = {
    DECEASED=make_cri_unitst_handlers(creatures.cri_unit.DECEASED),
    PET_OT={search_fn=overall_training_search},
    PET_AT={search_fn=assign_trainer_search},
+    WORK_DETAILS={search_fn=work_details_search},
    INTERROGATING={search_fn=interrogating_search},
    CONVICTING={search_fn=convicting_search},
 }
@ -213,6 +263,10 @@ local function get_key()
        return df.unit_list_mode_type[creatures.current_mode]
    elseif info.current_mode == df.info_interface_mode_type.ARTIFACTS then
        return df.artifacts_mode_type[objects.mode]
+    elseif info.current_mode == df.info_interface_mode_type.LABOR then
+        if info.labor.mode == df.labor_mode_type.WORK_DETAILS then
+            return 'WORK_DETAILS'
+        end
    end
 end

@ -225,8 +279,8 @@ local function check_context(self)
    end
 end

-local function do_search(self, text)
-    if not next(state) and text == '' then return end
+local function do_search(matches_filters_fn, text, force_full_search)
+    if not force_full_search and not next(state) and text == '' then return end
    -- the EditField state is guaranteed to be consistent with the current
    -- context since when clicking to switch tabs, onRenderBody is always called
    -- before this text_input callback, even if a key is pressed before the next
@ -238,11 +292,34 @@ local function do_search(self, text)
    local prev_text = ensure_key(state, key).prev_text
    -- some screens reset their contents between context switches; regardless
    -- a switch back to the context should results in an incremental search
-    local incremental = prev_text and text:startswith(prev_text)
-    HANDLERS[key].search_fn(state[key], text, incremental)
+    local incremental = not force_full_search and prev_text and text:startswith(prev_text)
+    HANDLERS[key].search_fn(matches_filters_fn, state[key], text, incremental)
    state[key].prev_text = text
 end

+local function on_update(self)
+    if self.overlay_onupdate_max_freq_seconds == 0 and
+        not dfhack.gui.matchFocusString('dwarfmode/Info', dfhack.gui.getDFViewscreen(true))
+    then
+        for k,v in pairs(state) do
+            local cleanup_fn = safe_index(HANDLERS, k, 'cleanup_fn')
+            if cleanup_fn then cleanup_fn(v) end
+        end
+        state = {}
+        self.subviews.search:setText('')
+        self.subviews.search:setFocus(false)
+        self.overlay_onupdate_max_freq_seconds = 60
+    end
+end
+
+local function on_input(self, clazz, keys)
+    if keys._MOUSE_R and self.subviews.search.focus then
+        self.subviews.search:setFocus(false)
+        return true
+    end
+    return clazz.super.onInput(self, keys)
+end
+
 -- ----------------------
 -- InfoOverlay
 --
@ -262,6 +339,7 @@ InfoOverlay.ATTRS{
        'dwarfmode/Info/ARTIFACTS/SYMBOLS',
        'dwarfmode/Info/ARTIFACTS/NAMED_OBJECTS',
        'dwarfmode/Info/ARTIFACTS/WRITTEN_CONTENT',
+        'dwarfmode/Info/LABOR/WORK_DETAILS',
    },
    hotspot=true,
    overlay_onupdate_max_freq_seconds=0,
@ -279,30 +357,15 @@ function InfoOverlay:init()
                    frame={l=1, t=0, r=1},
                    label_text="Search: ",
                    key='CUSTOM_ALT_S',
-                    on_change=curry(do_search, self),
+                    on_change=curry(do_search, DEFAULT_NIL),
                },
            },
        },
    }
 end

-local function cleanup()
-    for k,v in pairs(state) do
-        local cleanup_fn = safe_index(HANDLERS, k, 'cleanup_fn')
-        if cleanup_fn then cleanup_fn(v) end
-    end
-end
-
 function InfoOverlay:overlay_onupdate()
-    if next(state) and
-        not dfhack.gui.matchFocusString('dwarfmode/Info', dfhack.gui.getDFViewscreen(true))
-    then
-        cleanup()
-        state = {}
-        self.subviews.search:setText('')
-        self.subviews.search:setFocus(false)
-        self.overlay_onupdate_max_freq_seconds = 60
-    end
+    on_update(self)
 end

 local function resize_overlay(self)
@ -320,11 +383,12 @@ end

 local function get_panel_offsets()
    local tabs_in_two_rows = is_tabs_in_two_rows()
-    local is_objects = info.current_mode == df.info_interface_mode_type.ARTIFACTS
-    local l_offset = (not tabs_in_two_rows and is_objects) and 4 or 0
+    local shift_right = info.current_mode == df.info_interface_mode_type.ARTIFACTS or
+        info.current_mode == df.info_interface_mode_type.LABOR
+    local l_offset = (not tabs_in_two_rows and shift_right) and 4 or 0
    local t_offset = 1
    if tabs_in_two_rows then
-        t_offset = is_objects and 0 or 3
+        t_offset = shift_right and 0 or 3
    end
    return l_offset, t_offset
 end
@ -345,16 +409,19 @@ function InfoOverlay:onRenderBody(dc)
    if self:updateFrames() then
        self:updateLayout()
    end
+    if self.refresh_search then
+        self.refresh_search = nil
+        do_search(DEFAULT_NIL, self.subviews.search.text)
+    end
    self.overlay_onupdate_max_freq_seconds = 0
    InfoOverlay.super.onRenderBody(self, dc)
 end

 function InfoOverlay:onInput(keys)
-    if keys._MOUSE_R and self.subviews.search.focus then
-        self.subviews.search:setFocus(false)
-        return true
+    if keys._MOUSE_L and get_key() == 'WORK_DETAILS' then
+        self.refresh_search = true
    end
-    return InfoOverlay.super.onInput(self, keys)
+    return on_input(self, InfoOverlay, keys)
 end

 -- ----------------------
@ -367,8 +434,14 @@ InterrogationOverlay.ATTRS{
    default_enabled=true,
    viewscreens='dwarfmode/Info/JUSTICE',
    frame={w=27, h=9},
+    hotspot=true,
+    overlay_onupdate_max_freq_seconds=0,
 }

+function InterrogationOverlay:overlay_onupdate()
+    on_update(self)
+end
+
 local function is_interrogate_or_convict()
    local key = get_key()
    return key == 'INTERROGATING' or key == 'CONVICTING'
@ -388,39 +461,75 @@ function InterrogationOverlay:init()
                    frame={l=0, t=0, r=0},
                    label_text="Search: ",
                    key='CUSTOM_ALT_S',
-                    on_change=curry(do_search, self),
+                    on_change=curry(do_search, self:callback('matches_filters')),
+                },
+                widgets.ToggleHotkeyLabel{
+                    view_id='include_interviewed',
+                    frame={l=0, t=1, w=23},
+                    key='CUSTOM_SHIFT_I',
+                    label='Interviewed:',
+                    options={
+                        {label='Include', value=true, pen=COLOR_GREEN},
+                        {label='Exclude', value=false, pen=COLOR_RED},
+                    },
+                    visible=function() return justice.interrogating end,
+                    on_change=function()
+                        do_search(self:callback('matches_filters'), self.subviews.search.text, true)
+                    end,
                },
                widgets.CycleHotkeyLabel{
                    view_id='subset',
-                    frame={l=0, t=1, w=24},
+                    frame={l=0, t=2, w=28},
                    key='CUSTOM_SHIFT_F',
                    label='Show:',
                    options={
                        {label='All', value='all', pen=COLOR_GREEN},
-                        {label='Undead visitors', value='undead', pen=COLOR_RED},
+                        {label='Risky visitors', value='risky', pen=COLOR_RED},
                        {label='Other visitors', value='visitors', pen=COLOR_LIGHTRED},
                        {label='Residents', value='residents', pen=COLOR_YELLOW},
                        {label='Citizens', value='citizens', pen=COLOR_CYAN},
-                        {label='Animals', value='animals', pen=COLOR_MAGENTA},
-                        {label='Deceased', value='deceased', pen=COLOR_BLUE},
-                    },
-                },
-                widgets.ToggleHotkeyLabel{
-                    view_id='include_interviewed',
-                    frame={l=0, t=2, w=23},
-                    key='CUSTOM_SHIFT_I',
-                    label='Interviewed:',
-                    options={
-                        {label='Include', value=true, pen=COLOR_GREEN},
-                        {label='Exclude', value=false, pen=COLOR_RED},
+                        {label='Animals', value='animals', pen=COLOR_BLUE},
+                        {label='Deceased or missing', value='deceased', pen=COLOR_MAGENTA},
+                        {label='Others', value='others', pen=COLOR_GRAY},
                    },
-                    initial_option=true,
+                    on_change=function()
+                        do_search(self:callback('matches_filters'), self.subviews.search.text, true)
+                    end,
                },
            },
        },
    }
 end

+local function is_risky(unit)
+    return false
+end
+
+function InterrogationOverlay:matches_filters(unit, idx)
+    if justice.interrogating then
+        local include_interviewed = self.subviews.include_interviewed:getOptionValue()
+        if not include_interviewed and justice.interrogation_list_flag[idx] == 2 then
+            return false
+        end
+    end
+    local subset = self.subviews.subset:getOptionValue()
+    if subset == 'all' then
+        return true
+    elseif dfhack.units.isDead(unit) or not dfhack.units.isActive(unit) then
+        return subset == 'deceased'
+    elseif dfhack.units.isVisiting(unit) then
+        local risky = is_risky(unit)
+        return (subset == 'risky' and risky) or (subset == 'visitors' and not risky)
+    elseif dfhack.units.isAnimal(unit) then
+        return subset == 'animals'
+    elseif dfhack.units.isCitizen(unit) then
+        return subset == 'citizens'
+    elseif dfhack.units.isOwnGroup(unit) then
+        return subset == 'residents'
+    end
+    return subset == 'others'
+end
+
 function InterrogationOverlay:render(dc)
    local sw = dfhack.screen.getWindowSize()
    local info_panel_border = 31 -- from edges of panel to screen edges
@ -447,16 +556,16 @@ end
 function InterrogationOverlay:onRenderBody(dc)
    if next(state) then
        check_context(self)
+    else
+        self.subviews.include_interviewed:setOption(true, false)
+        self.subviews.subset:setOption('all')
    end
+    self.overlay_onupdate_max_freq_seconds = 0
    InterrogationOverlay.super.onRenderBody(self, dc)
 end

 function InterrogationOverlay:onInput(keys)
-    if keys._MOUSE_R and self.subviews.search.focus then
-        self.subviews.search:setFocus(false)
-        return true
-    end
-    return InterrogationOverlay.super.onInput(self, keys)
+    return on_input(self, InterrogationOverlay, keys)
 end

 return _ENV