From 09e3ed427a2967682d54cc3a607028402c658ef9 Mon Sep 17 00:00:00 2001
From: Myk Taylor <myk002@yahoo.com>
Date: Tue, 10 Oct 2023 05:03:59 -0700
Subject: [PATCH] add search functionality for burrows assignment screen

---
 docs/changelog.txt           |  1 +
 plugins/lua/sort.lua         |  9 +++--
 plugins/lua/sort/burrows.lua | 70 ++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 5 deletions(-)
 create mode 100644 plugins/lua/sort/burrows.lua

diff --git a/docs/changelog.txt b/docs/changelog.txt
index 589952305..f79cbdd28 100644
--- a/docs/changelog.txt
+++ b/docs/changelog.txt
@@ -59,6 +59,7 @@ Template for new versions:
 - `logistics`: ``automelt`` now optionally supports melting masterworks; feature accessible from `stockpiles` overlay
 - `sort`: new search widgets for Info panel tabs, including all "Creatures" subtabs, all "Objects" subtabs, "Tasks", the "Work details" subtab under "Labor", and the "Interrogate" and "Convict" screens under "Justice"
 - `sort`: new search widgets for location selection screen (when you're choosing what kind of guildhall or temple to dedicate)
+- `sort`: new search widgets for burrow assignment screen
 
 ## Fixes
 - `zone`: don't show animal assignment link for dungeon cages/restraints
diff --git a/plugins/lua/sort.lua b/plugins/lua/sort.lua
index ce9eb1f9e..3bc97cb74 100644
--- a/plugins/lua/sort.lua
+++ b/plugins/lua/sort.lua
@@ -1,8 +1,6 @@
 local _ENV = mkmodule('plugins.sort')
 
-local info = require('plugins.sort.info')
 local gui = require('gui')
-local locationselector = require('plugins.sort.locationselector')
 local overlay = require('plugins.overlay')
 local setbelief = reqscript('modtools/set-belief')
 local textures = require('gui.textures')
@@ -1287,9 +1285,10 @@ end
 OVERLAY_WIDGETS = {
     squad_assignment=SquadAssignmentOverlay,
     squad_annotation=SquadAnnotationOverlay,
-    info=info.InfoOverlay,
-    interrogation=info.InterrogationOverlay,
-    location_selector=locationselector.LocationSelectorOverlay,
+    info=require('plugins.sort.info').InfoOverlay,
+    interrogation=require('plugins.sort.info').InterrogationOverlay,
+    location_selector=require('plugins.sort.locationselector').LocationSelectorOverlay,
+    burrows=require('plugins.sort.burrows').BurrowOverlay,
 }
 
 dfhack.onStateChange[GLOBAL_KEY] = function(sc)
diff --git a/plugins/lua/sort/burrows.lua b/plugins/lua/sort/burrows.lua
new file mode 100644
index 000000000..fe2b8ec8e
--- /dev/null
+++ b/plugins/lua/sort/burrows.lua
@@ -0,0 +1,70 @@
+local _ENV = mkmodule('plugins.sort.burrows')
+
+local sortoverlay = require('plugins.sort.sortoverlay')
+local widgets = require('gui.widgets')
+
+local unit_selector = df.global.game.main_interface.unit_selector
+
+-- ----------------------
+-- BurrowOverlay
+--
+
+BurrowOverlay = defclass(BurrowOverlay, sortoverlay.SortOverlay)
+BurrowOverlay.ATTRS{
+    default_pos={x=62, y=6},
+    viewscreens='dwarfmode/UnitSelector/BURROW_ASSIGNMENT',
+    frame={w=26, h=1},
+}
+
+local function get_unit_id_search_key(unit_id)
+    local unit = df.unit.find(unit_id)
+    if not unit then return end
+    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
+
+function BurrowOverlay:init()
+    self:addviews{
+        widgets.BannerPanel{
+            frame={l=0, t=0, r=0, h=1},
+            subviews={
+                widgets.EditField{
+                    view_id='search',
+                    frame={l=1, t=0, r=1},
+                    label_text="Search: ",
+                    key='CUSTOM_ALT_S',
+                    on_change=function(text) self:do_search(text) end,
+                },
+            },
+        },
+    }
+
+    self:register_handler('BURROW', unit_selector.unid,
+        curry(sortoverlay.flags_vector_search, {get_search_key_fn=get_unit_id_search_key},
+            unit_selector.selected))
+end
+
+function BurrowOverlay:get_key()
+    if unit_selector.context == df.unit_selector_context_type.BURROW_ASSIGNMENT then
+        return 'BURROW'
+    end
+end
+
+function BurrowOverlay:onRenderBody(dc)
+    BurrowOverlay.super.onRenderBody(self, dc)
+    if self.refresh_search then
+        self.refresh_search = nil
+        self:do_search(self.subviews.search.text)
+    end
+end
+
+function BurrowOverlay:onInput(keys)
+    if keys._MOUSE_L then
+        self.refresh_search = true
+    end
+    return BurrowOverlay.super.onInput(self, keys)
+end
+
+return _ENV