From e196ceb8f9ddb27bfb995931137b810d7b78c630 Mon Sep 17 00:00:00 2001
From: Myk Taylor <myk002@yahoo.com>
Date: Sat, 4 Nov 2023 18:59:20 -0700
Subject: [PATCH] support search for work animal assignment

---
 docs/changelog.txt        |  1 +
 plugins/lua/sort/info.lua | 47 +++++++++++++++++++++++++++++++++------
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/docs/changelog.txt b/docs/changelog.txt
index 7a13ba6f4..a6533ef82 100644
--- a/docs/changelog.txt
+++ b/docs/changelog.txt
@@ -66,6 +66,7 @@ Template for new versions:
 ## Misc Improvements
 - `buildingplan`: display how many items are available on the planner panel
 - `sort`: allow searching by profession on the squad assignment page
+- `sort`: add search for work animal assignment screen; allow filtering by miltary/civilian
 - `dreamfort`: Inside+ and Clearcutting burrows now automatically created and managed
 
 ## Documentation
diff --git a/plugins/lua/sort/info.lua b/plugins/lua/sort/info.lua
index a48b1d986..f9233bf6c 100644
--- a/plugins/lua/sort/info.lua
+++ b/plugins/lua/sort/info.lua
@@ -109,11 +109,6 @@ local function get_race_name(raw_id)
     return raw.name[1]
 end
 
-local function get_trainer_search_key(unit)
-    if not unit then return end
-    return ('%s %s'):format(dfhack.TranslateName(unit.name), dfhack.units.getProfessionName(unit))
-end
-
 -- get name in both dwarvish and English
 local function get_artifact_search_key(artifact)
     return ('%s %s'):format(dfhack.TranslateName(artifact.name), dfhack.TranslateName(artifact.name, true))
@@ -166,7 +161,7 @@ InfoOverlay = defclass(InfoOverlay, sortoverlay.SortOverlay)
 InfoOverlay.ATTRS{
     default_pos={x=64, y=8},
     viewscreens='dwarfmode/Info',
-    frame={w=40, h=4},
+    frame={w=40, h=5},
 }
 
 function InfoOverlay:init()
@@ -185,6 +180,25 @@ function InfoOverlay:init()
                 },
             },
         },
+        widgets.BannerPanel{
+            view_id='filter_panel',
+            frame={l=0, t=1, r=0, h=1},
+            visible=function() return self:get_key() == 'PET_WA' end,
+            subviews={
+                widgets.CycleHotkeyLabel{
+                    view_id='subset',
+                    frame={l=1, t=0},
+                    key='CUSTOM_SHIFT_F',
+                    label='Show:',
+                    options={
+                        {label='All', value='all', pen=COLOR_GREEN},
+                        {label='Military', value='military', pen=COLOR_YELLOW},
+                        {label='Civilians', value='civilian', pen=COLOR_CYAN},
+                    },
+                    on_change=function() self:do_search(self.subviews.search.text, true) end,
+                },
+            },
+        },
     }
 
     local CRI_UNIT_VECS = {
@@ -209,7 +223,12 @@ function InfoOverlay:init()
     self:register_handler('PET_OT', creatures.atk_index,
         curry(sortoverlay.single_vector_search, {get_search_key_fn=get_race_name}))
     self:register_handler('PET_AT', creatures.trainer,
-        curry(sortoverlay.single_vector_search, {get_search_key_fn=get_trainer_search_key}))
+        curry(sortoverlay.single_vector_search, {get_search_key_fn=sortoverlay.get_unit_search_key}))
+    self:register_handler('PET_WA', creatures.work_animal_recipient,
+        curry(sortoverlay.single_vector_search, {
+            get_search_key_fn=sortoverlay.get_unit_search_key,
+            matches_filters_fn=self:callback('matches_filters'),
+        }))
     self:register_handler('WORK_DETAILS', work_details.assignable_unit, work_details_search)
 
     for idx,name in ipairs(df.artifacts_mode_type) do
@@ -227,6 +246,8 @@ function InfoOverlay:get_key()
                 return 'PET_OT'
             elseif creatures.adding_trainer then
                 return 'PET_AT'
+            elseif creatures.assign_work_animal then
+                return 'PET_WA'
             end
         end
         return df.unit_list_mode_type[creatures.current_mode]
@@ -275,6 +296,8 @@ function InfoOverlay:updateFrames()
     local frame = self.subviews.panel.frame
     if frame.l == l and frame.t == t then return ret end
     frame.l, frame.t = l, t
+    local frame2 = self.subviews.filter_panel.frame
+    frame2.l, frame2.t = l, t + 1
     return true
 end
 
@@ -296,6 +319,16 @@ function InfoOverlay:onInput(keys)
     return InfoOverlay.super.onInput(self, keys)
 end
 
+function InfoOverlay:matches_filters(unit)
+    local subset = self.subviews.subset:getOptionValue()
+    if subset == 'all' then
+        return true
+    elseif unit.military.squad_id == -1 then
+        return subset == 'civilian'
+    end
+    return subset == 'military'
+end
+
 -- ----------------------
 -- CandidatesOverlay
 --