From 60787dec3143bad7e02425efce9f4322df9dabe8 Mon Sep 17 00:00:00 2001
From: Myk Taylor <myk002@yahoo.com>
Date: Thu, 4 Jan 2024 03:58:04 -0800
Subject: [PATCH] clean up occupations vector when retiring loc

also close the details dialog immediately so the player doesn't have a
chance to reassign occupations
---
 plugins/lua/zone.lua | 58 +++++++++++++++++++++++---------------------
 1 file changed, 31 insertions(+), 27 deletions(-)

diff --git a/plugins/lua/zone.lua b/plugins/lua/zone.lua
index 8e16ec219..802396028 100644
--- a/plugins/lua/zone.lua
+++ b/plugins/lua/zone.lua
@@ -1103,13 +1103,12 @@ RetireLocationOverlay.ATTRS{
     default_pos={x=-39,y=6},
     default_enabled=true,
     viewscreens='dwarfmode/LocationDetails',
-    frame={w=25, h=4},
+    frame={w=25, h=3},
 }
 
 function RetireLocationOverlay:init()
     self:addviews{
         widgets.Panel{
-            frame={l=0, t=0, r=0, h=3},
             frame_background=gui.CLEAR_PEN,
             frame_style=gui.FRAME_MEDIUM,
             visible=location_details_is_on_top,
@@ -1118,36 +1117,49 @@ function RetireLocationOverlay:init()
                     frame={t=0, l=0},
                     label='Retire location',
                     key='CUSTOM_CTRL_D',
-                    on_activate=self:callback('retire'),
+                    on_activate=self:callback('confirm_retire'),
                 },
             },
         },
-        widgets.Label{
-            frame={l=1, b=0},
-            text='LOCATION RETIRED',
-            text_pen=COLOR_RED,
-            visible=function() return mi.location_details.selected_ab.flags.DOES_NOT_EXIST end
-        },
     }
 end
 
-function RetireLocationOverlay:retire()
+local function retire(location)
+    location.flags.DOES_NOT_EXIST = true
+    for idx, loc in ipairs(mi.location_selector.valid_ab) do
+        if loc.id == location.id then
+            mi.location_selector.valid_ab:erase(idx)
+            break
+        end
+    end
+    local occupations = df.global.world.occupations.all
+    for idx = #occupations-1,0,-1 do
+        local occupation = occupations[idx]
+        if occupation.site_id == location.site_id and
+            occupation.location_id == location.id
+        then
+            occupations:erase(idx)
+        end
+    end
+    mi.location_details.open = false
+end
+
+function RetireLocationOverlay:confirm_retire()
     local details = mi.location_details
     local location = details.selected_ab
-    local has_occupations, has_zones = false, #location.contents.building_ids ~= 0
+    local num_occupations, num_zones = 0, #location.contents.building_ids
     for _, occupation in ipairs(location.occupations) do
         if occupation.histfig_id ~= -1 then
-            has_occupations = true
-            break
+            num_occupations = num_occupations + 1
         end
     end
-    if has_occupations or has_zones then
+    if num_occupations + num_zones > 0 then
         local messages = {'Cannot retire location! Please:', ''}
-        if has_occupations then
-            table.insert(messages, '- unassign location occupations')
+        if num_occupations > 0 then
+            table.insert(messages, ('- unassign %d location occupation(s)'):format(num_occupations))
         end
-        if has_zones then
-            table.insert(messages, ('- detach this location from %d zone(s)'):format(#location.contents.building_ids))
+        if num_zones > 0 then
+            table.insert(messages, ('- detach this location from %d zone(s)'):format(num_zones))
         end
         table.insert(messages, '')
         table.insert(messages, 'and try again')
@@ -1157,15 +1169,7 @@ function RetireLocationOverlay:retire()
     dialogs.showYesNoPrompt('Confirm retire location',
         'Are you sure you want to retire this location?'..NEWLINE..'You won\'t be able to use it again.',
         COLOR_WHITE,
-        function()
-            location.flags.DOES_NOT_EXIST = true
-            for idx, loc in ipairs(mi.location_selector.valid_ab) do
-                if loc.id == location.id then
-                    mi.location_selector.valid_ab:erase(idx)
-                    break
-                end
-            end
-        end)
+        curry(retire, location))
 end
 
 OVERLAY_WIDGETS = {