local _ENV = mkmodule('plugins.zone') local dialogs = require('gui.dialogs') local gui = require('gui') local overlay = require('plugins.overlay') local utils = require('utils') local widgets = require('gui.widgets') local CH_UP = string.char(30) local CH_DN = string.char(31) local CH_MALE = string.char(11) local CH_FEMALE = string.char(12) local CH_NEUTER = '?' local DISPOSITION = { NONE={label='Unknown', value=0}, PET={label='Pet', value=1}, TAME={label='Domesticated', value=2}, TRAINED={label='Partially trained', value=3}, WILD_TRAINABLE={label='Wild (trainable)', value=4}, WILD_UNTRAINABLE={label='Wild (untrainable)', value=5}, HOSTILE={label='Hostile', value=6}, } local DISPOSITION_REVMAP = {} for k, v in pairs(DISPOSITION) do DISPOSITION_REVMAP[v.value] = k end -- ------------------- -- AssignAnimal -- local STATUS_COL_WIDTH = 18 local DISPOSITION_COL_WIDTH = 18 local GENDER_COL_WIDTH = 6 local SLIDER_LABEL_WIDTH = math.max(STATUS_COL_WIDTH, DISPOSITION_COL_WIDTH) + 4 local SLIDER_WIDTH = 48 AssignAnimal = defclass(AssignAnimal, widgets.Window) AssignAnimal.ATTRS { frame_title='Animal and prisoner assignment', frame={w=6+SLIDER_WIDTH*2, h=47}, resizable=true, resize_min={h=27}, target_name=DEFAULT_NIL, status=DEFAULT_NIL, status_revmap=DEFAULT_NIL, get_status=DEFAULT_NIL, get_allow_vermin=DEFAULT_NIL, get_multi_select=DEFAULT_NIL, attach=DEFAULT_NIL, initial_min_disposition=DISPOSITION.PET.value, } local function sort_noop(a, b) -- this function is used as a marker and never actually gets called error('sort_noop should not be called') end local function sort_by_race_desc(a, b) if a.data.race == b.data.race then return a.data.gender < b.data.gender end return a.data.race < b.data.race end local function sort_by_race_asc(a, b) if a.data.race == b.data.race then return a.data.gender < b.data.gender end return a.data.race > b.data.race end local function sort_by_name_desc(a, b) if a.search_key == b.search_key then return sort_by_race_desc(a, b) end return a.search_key < b.search_key end local function sort_by_name_asc(a, b) if a.search_key == b.search_key then return sort_by_race_desc(a, b) end return a.search_key > b.search_key end local function sort_by_gender_desc(a, b) if a.data.gender == b.data.gender then return sort_by_race_desc(a, b) end return a.data.gender < b.data.gender end local function sort_by_gender_asc(a, b) if a.data.gender == b.data.gender then return sort_by_race_desc(a, b) end return a.data.gender > b.data.gender end local function sort_by_disposition_desc(a, b) if a.data.disposition == b.data.disposition then return sort_by_race_desc(a, b) end return a.data.disposition < b.data.disposition end local function sort_by_disposition_asc(a, b) if a.data.disposition == b.data.disposition then return sort_by_race_desc(a, b) end return a.data.disposition > b.data.disposition end local function sort_by_status_desc(a, b) if a.data.status == b.data.status then return sort_by_race_desc(a, b) end return a.data.status < b.data.status end local function sort_by_status_asc(a, b) if a.data.status == b.data.status then return sort_by_race_desc(a, b) end return a.data.status > b.data.status end function AssignAnimal:init() local status_options = {} for k, v in ipairs(self.status_revmap) do status_options[k] = self.status[v] end local disposition_options = {} for k, v in ipairs(DISPOSITION_REVMAP) do disposition_options[k] = DISPOSITION[v] end local can_assign_pets = self.initial_min_disposition == DISPOSITION.PET.value self:addviews{ widgets.CycleHotkeyLabel{ view_id='sort', frame={l=0, t=0, w=26}, label='Sort by:', key='CUSTOM_SHIFT_S', options={ {label='status'..CH_DN, value=sort_by_status_desc}, {label='status'..CH_UP, value=sort_by_status_asc}, {label='disposition'..CH_DN, value=sort_by_disposition_desc}, {label='disposition'..CH_UP, value=sort_by_disposition_asc}, {label='gender'..CH_DN, value=sort_by_gender_desc}, {label='gender'..CH_UP, value=sort_by_gender_asc}, {label='race'..CH_DN, value=sort_by_race_desc}, {label='race'..CH_UP, value=sort_by_race_asc}, {label='name'..CH_DN, value=sort_by_name_desc}, {label='name'..CH_UP, value=sort_by_name_asc}, }, initial_option=sort_by_status_desc, on_change=self:callback('refresh_list', 'sort'), }, widgets.EditField{ view_id='search', frame={l=35, t=0}, label_text='Search: ', on_char=function(ch) return ch:match('[%l -]') end, }, widgets.Panel{ frame={t=2, l=0, w=SLIDER_WIDTH, h=4}, subviews={ widgets.CycleHotkeyLabel{ view_id='min_status', frame={l=0, t=0, w=SLIDER_LABEL_WIDTH}, label='Min status:', label_below=true, key_back='CUSTOM_SHIFT_Z', key='CUSTOM_SHIFT_X', options=status_options, initial_option=1, on_change=function(val) if self.subviews.max_status:getOptionValue() < val then self.subviews.max_status:setOption(val) end self:refresh_list() end, }, widgets.CycleHotkeyLabel{ view_id='max_status', frame={r=1, t=0, w=SLIDER_LABEL_WIDTH}, label='Max status:', label_below=true, key_back='CUSTOM_SHIFT_Q', key='CUSTOM_SHIFT_W', options=status_options, initial_option=#self.status_revmap, on_change=function(val) if self.subviews.min_status:getOptionValue() > val then self.subviews.min_status:setOption(val) end self:refresh_list() end, }, widgets.RangeSlider{ frame={l=0, t=3}, num_stops=#self.status_revmap, get_left_idx_fn=function() return self.subviews.min_status:getOptionValue() end, get_right_idx_fn=function() return self.subviews.max_status:getOptionValue() end, on_left_change=function(idx) self.subviews.min_status:setOption(idx, true) end, on_right_change=function(idx) self.subviews.max_status:setOption(idx, true) end, }, }, }, widgets.Panel{ frame={t=7, l=0, w=SLIDER_WIDTH, h=4}, subviews={ widgets.CycleHotkeyLabel{ view_id='min_disposition', frame={l=0, t=0, w=SLIDER_LABEL_WIDTH}, label='Min disposition:', label_below=true, key_back='CUSTOM_SHIFT_C', key='CUSTOM_SHIFT_V', options=disposition_options, initial_option=self.initial_min_disposition, on_change=function(val) if self.subviews.max_disposition:getOptionValue() < val then self.subviews.max_disposition:setOption(val) end self:refresh_list() end, }, widgets.CycleHotkeyLabel{ view_id='max_disposition', frame={r=1, t=0, w=SLIDER_LABEL_WIDTH}, label='Max disposition:', label_below=true, key_back='CUSTOM_SHIFT_E', key='CUSTOM_SHIFT_R', options=disposition_options, initial_option=DISPOSITION.HOSTILE.value, on_change=function(val) if self.subviews.min_disposition:getOptionValue() > val then self.subviews.min_disposition:setOption(val) end self:refresh_list() end, }, widgets.RangeSlider{ frame={l=0, t=3}, num_stops=6, get_left_idx_fn=function() return self.subviews.min_disposition:getOptionValue() end, get_right_idx_fn=function() return self.subviews.max_disposition:getOptionValue() end, on_left_change=function(idx) self.subviews.min_disposition:setOption(idx, true) end, on_right_change=function(idx) self.subviews.max_disposition:setOption(idx, true) end, }, }, }, widgets.Panel{ frame={t=3, l=SLIDER_WIDTH+2, r=0, h=3}, subviews={ widgets.CycleHotkeyLabel{ view_id='egg', frame={l=0, t=0, w=23}, key_back='CUSTOM_SHIFT_B', key='CUSTOM_SHIFT_N', label='Egg layers:', options={ {label='Include', value='include', pen=COLOR_GREEN}, {label='Only', value='only', pen=COLOR_YELLOW}, {label='Exclude', value='exclude', pen=COLOR_RED}, }, initial_option='include', on_change=function() self:refresh_list() end, }, widgets.CycleHotkeyLabel{ view_id='graze', frame={l=0, t=2, w=20}, key_back='CUSTOM_SHIFT_T', key='CUSTOM_SHIFT_Y', label='Grazers:', options={ {label='Include', value='include', pen=COLOR_GREEN}, {label='Only', value='only', pen=COLOR_YELLOW}, {label='Exclude', value='exclude', pen=COLOR_RED}, }, initial_option='include', on_change=function() self:refresh_list() end, }, }, }, widgets.Panel{ view_id='list_panel', frame={t=12, l=0, r=0, b=4+(can_assign_pets and 0 or 1)}, subviews={ widgets.CycleHotkeyLabel{ view_id='sort_status', frame={t=0, l=0, w=7}, options={ {label='status', value=sort_noop}, {label='status'..CH_DN, value=sort_by_status_desc}, {label='status'..CH_UP, value=sort_by_status_asc}, }, initial_option=sort_by_status_desc, option_gap=0, on_change=self:callback('refresh_list', 'sort_status'), }, widgets.CycleHotkeyLabel{ view_id='sort_disposition', frame={t=0, l=STATUS_COL_WIDTH+2, w=12}, options={ {label='disposition', value=sort_noop}, {label='disposition'..CH_DN, value=sort_by_disposition_desc}, {label='disposition'..CH_UP, value=sort_by_disposition_asc}, }, option_gap=0, on_change=self:callback('refresh_list', 'sort_disposition'), }, widgets.CycleHotkeyLabel{ view_id='sort_gender', frame={t=0, l=STATUS_COL_WIDTH+2+DISPOSITION_COL_WIDTH+2, w=7}, options={ {label='gender', value=sort_noop}, {label='gender'..CH_DN, value=sort_by_gender_desc}, {label='gender'..CH_UP, value=sort_by_gender_asc}, }, option_gap=0, on_change=self:callback('refresh_list', 'sort_gender'), }, widgets.CycleHotkeyLabel{ view_id='sort_race', frame={t=0, l=STATUS_COL_WIDTH+2+DISPOSITION_COL_WIDTH+2+GENDER_COL_WIDTH+2, w=5}, options={ {label='race', value=sort_noop}, {label='race'..CH_DN, value=sort_by_race_desc}, {label='race'..CH_UP, value=sort_by_race_asc}, }, option_gap=0, on_change=self:callback('refresh_list', 'sort_race'), }, widgets.CycleHotkeyLabel{ view_id='sort_name', frame={t=0, l=STATUS_COL_WIDTH+2+DISPOSITION_COL_WIDTH+2+GENDER_COL_WIDTH+2+7, w=5}, options={ {label='name', value=sort_noop}, {label='name'..CH_DN, value=sort_by_name_desc}, {label='name'..CH_UP, value=sort_by_name_asc}, }, option_gap=0, on_change=self:callback('refresh_list', 'sort_name'), }, widgets.FilteredList{ view_id='list', frame={l=0, t=2, r=0, b=0}, on_submit=self:callback('toggle_item'), on_submit2=self:callback('toggle_range'), on_select=self:callback('select_item'), }, } }, widgets.HotkeyLabel{ frame={l=0, b=2+(can_assign_pets and 0 or 1)}, label='Assign all/none', key='CUSTOM_CTRL_A', on_activate=self:callback('toggle_visible'), visible=self.get_multi_select, auto_width=true, }, widgets.WrappedLabel{ frame={b=0, l=0, r=0}, text_to_wrap=function() return 'Click to assign/unassign to ' .. self.target_name .. '.' .. (self.get_multi_select() and ' Shift click to assign/unassign a range.' or '') .. (not can_assign_pets and '\nNote that pets cannot be assigned to cages or restraints.' or '') end, }, } -- replace the FilteredList's built-in EditField with our own self.subviews.list.list.frame.t = 0 self.subviews.list.edit.visible = false self.subviews.list.edit = self.subviews.search self.subviews.search.on_change = self.subviews.list:callback('onFilterChange') self.subviews.list:setChoices(self:get_choices()) end function AssignAnimal:refresh_list(sort_widget, sort_fn) sort_widget = sort_widget or 'sort' sort_fn = sort_fn or self.subviews.sort:getOptionValue() if sort_fn == sort_noop then self.subviews[sort_widget]:cycle() return end for _,widget_name in ipairs{'sort', 'sort_status', 'sort_disposition', 'sort_gender', 'sort_race', 'sort_name'} do self.subviews[widget_name]:setOption(sort_fn) end local list = self.subviews.list local saved_filter = list:getFilter() list:setFilter('') list:setChoices(self:get_choices(), list:getSelected()) list:setFilter(saved_filter) end function add_words(words, str) for word in dfhack.toSearchNormalized(str):gmatch("[%w-]+") do table.insert(words, word:lower()) end end function make_search_key(desc, race_raw) local words = {} add_words(words, desc) if race_raw then add_words(words, race_raw.name[0]) end return table.concat(words, ' ') end function AssignAnimal:make_choice_text(data) local gender_ch = CH_NEUTER if data.gender == df.pronoun_type.she then gender_ch = CH_FEMALE elseif data.gender == df.pronoun_type.he then gender_ch = CH_MALE end return { {width=STATUS_COL_WIDTH, text=function() return self.status[self.status_revmap[data.status]].label end}, {gap=2, width=DISPOSITION_COL_WIDTH, text=function() return DISPOSITION[DISPOSITION_REVMAP[data.disposition]].label end}, {gap=2, width=GENDER_COL_WIDTH, text=gender_ch}, {gap=2, text=data.desc}, } end local function get_general_ref(unit_or_vermin, ref_type) local is_unit = df.unit:is_instance(unit_or_vermin) return dfhack[is_unit and 'units' or 'items'].getGeneralRef(unit_or_vermin, ref_type) end local function get_cage_ref(unit_or_vermin) return get_general_ref(unit_or_vermin, df.general_ref_type.CONTAINED_IN_ITEM) end local function get_built_cage(item_cage) if not item_cage then return end local built_cage_ref = dfhack.items.getGeneralRef(item_cage, df.general_ref_type.BUILDING_HOLDER) if not built_cage_ref then return end local built_cage = df.building.find(built_cage_ref.building_id) if not built_cage then return end if built_cage:getType() == df.building_type.Cage then return built_cage end end local function get_bld_assignments() local assignments = {} for _,cage in ipairs(df.global.world.buildings.other.CAGE) do for _,unit_id in ipairs(cage.assigned_units) do assignments[unit_id] = cage end end for _,chain in ipairs(df.global.world.buildings.other.CHAIN) do if chain.assigned then assignments[chain.assigned.id] = chain end end return assignments end local function get_unit_disposition(unit) local disposition = DISPOSITION.NONE if dfhack.units.isInvader(unit) or dfhack.units.isOpposedToLife(unit) then disposition = DISPOSITION.HOSTILE elseif dfhack.units.isPet(unit) then disposition = DISPOSITION.PET elseif dfhack.units.isDomesticated(unit) then disposition = DISPOSITION.TAME elseif dfhack.units.isTame(unit) then disposition = DISPOSITION.TRAINED elseif dfhack.units.isTamable(unit) then disposition = DISPOSITION.WILD_TRAINABLE else disposition = DISPOSITION.WILD_UNTRAINABLE end return disposition.value end local function get_item_disposition(item) if dfhack.units.casteFlagSet(item.race, item.caste, df.caste_raw_flags.OPPOSED_TO_LIFE) then return DISPOSITION.HOSTILE.value end if df.item_petst:is_instance(item) then if item.owner_id > -1 then return DISPOSITION.PET.value end return DISPOSITION.TAME.value end if dfhack.units.casteFlagSet(item.race, item.caste, df.caste_raw_flags.PET) or dfhack.units.casteFlagSet(item.race, item.caste, df.caste_raw_flags.PET_EXOTIC) then return DISPOSITION.WILD_TRAINABLE.value end return DISPOSITION.WILD_UNTRAINABLE.value end local function is_assignable_unit(unit) return dfhack.units.isActive(unit) and ((dfhack.units.isAnimal(unit) and dfhack.units.isOwnCiv(unit)) or get_cage_ref(unit)) and not dfhack.units.isDead(unit) and not dfhack.units.isMerchant(unit) and not dfhack.units.isForest(unit) end local function is_assignable_item(item) -- all vermin/small pets are assignable return true end local function get_vermin_desc(vermin, raw) if not raw then return 'Unknown vermin' end if vermin.stack_size > 1 then return ('%s [%d]'):format(raw.name[1], vermin.stack_size) end return ('%s'):format(raw.name[0]) end local function get_small_pet_desc(raw) if not raw then return 'Unknown small pet' end return ('tame %s'):format(raw.name[0]) end function AssignAnimal:cache_choices() if self.choices then return self.choices end local bld_assignments = get_bld_assignments() local choices = {} for _, unit in ipairs(df.global.world.units.active) do if not is_assignable_unit(unit) then goto continue end local raw = df.creature_raw.find(unit.race) local data = { unit=unit, desc=dfhack.units.getReadableName(unit), gender=unit.sex, race=raw and raw.creature_id or '', status=self.get_status(unit, bld_assignments), disposition=get_unit_disposition(unit), egg=dfhack.units.isEggLayerRace(unit), graze=dfhack.units.isGrazer(unit), } local choice = { search_key=make_search_key(data.desc, raw), data=data, text=self:make_choice_text(data), } table.insert(choices, choice) ::continue:: end for _, vermin in ipairs(df.global.world.items.other.VERMIN) do if not is_assignable_item(vermin) then goto continue end local raw = df.creature_raw.find(vermin.race) local data = { vermin=vermin, desc=get_vermin_desc(vermin, raw), gender=df.pronoun_type.it, race=raw and raw.creature_id or '', status=self.get_status(vermin, bld_assignments), disposition=get_item_disposition(vermin), } local choice = { search_key=make_search_key(data.desc, raw), data=data, text=self:make_choice_text(data), } table.insert(choices, choice) ::continue:: end for _, small_pet in ipairs(df.global.world.items.other.PET) do if not is_assignable_item(small_pet) then goto continue end local raw = df.creature_raw.find(small_pet.race) local data = { vermin=small_pet, desc=get_small_pet_desc(raw), gender=df.pronoun_type.it, race=raw and raw.creature_id or '', status=self.get_status(small_pet, bld_assignments), disposition=get_item_disposition(small_pet), } local choice = { search_key=make_search_key(data.desc, raw), data=data, text=self:make_choice_text(data), } table.insert(choices, choice) ::continue:: end self.choices = choices return choices end function AssignAnimal:get_choices() local raw_choices = self:cache_choices() local show_vermin = self.get_allow_vermin() local min_status = self.subviews.min_status:getOptionValue() local max_status = self.subviews.max_status:getOptionValue() local min_disposition = self.subviews.min_disposition:getOptionValue() local max_disposition = self.subviews.max_disposition:getOptionValue() local egg = self.subviews.egg:getOptionValue() local graze = self.subviews.graze:getOptionValue() local choices = {} for _,choice in ipairs(raw_choices) do local data = choice.data if not show_vermin and data.vermin then goto continue end if min_status > data.status then goto continue end if max_status < data.status then goto continue end if min_disposition > data.disposition then goto continue end if max_disposition < data.disposition then goto continue end if egg == 'only' and not data.egg then goto continue end if egg == 'exclude' and data.egg then goto continue end if graze == 'only' and not data.graze then goto continue end if graze == 'exclude' and data.graze then goto continue end table.insert(choices, choice) ::continue:: end table.sort(choices, self.subviews.sort:getOptionValue()) return choices end local function get_bld_assigned_vec(bld, unit_or_vermin) if not bld then return end return df.unit:is_instance(unit_or_vermin) and bld.assigned_units or bld.assigned_items end local function get_assigned_unit_or_vermin_idx(bld, unit_or_vermin, vec) vec = vec or get_bld_assigned_vec(bld, unit_or_vermin) if not vec then return end for assigned_idx, assigned_id in ipairs(vec) do if assigned_id == unit_or_vermin.id then return assigned_idx end end end local function unassign_unit_or_vermin(bld, unit_or_vermin) if not bld then return end if df.building_chainst:is_instance(bld) then if bld.assigned == unit_or_vermin then bld.assigned = nil end return end local vec = get_bld_assigned_vec(bld, unit_or_vermin) local idx = get_assigned_unit_or_vermin_idx(bld, unit_or_vermin, vec) if vec and idx then vec:erase(idx) end end local function detach_unit_or_vermin(unit_or_vermin, bld_assignments) for idx = #unit_or_vermin.general_refs-1, 0, -1 do local ref = unit_or_vermin.general_refs[idx] if df.general_ref_building_civzone_assignedst:is_instance(ref) then unassign_unit_or_vermin(df.building.find(ref.building_id), unit_or_vermin) unit_or_vermin.general_refs:erase(idx) ref:delete() elseif df.general_ref_contained_in_itemst:is_instance(ref) then local built_cage = get_built_cage(df.item.find(ref.item_id)) if built_cage and built_cage:getType() == df.building_type.Cage then unassign_unit_or_vermin(built_cage, unit_or_vermin) -- unit's general ref will be removed when the unit is released from the cage end end end bld_assignments = bld_assignments or get_bld_assignments() if bld_assignments[unit_or_vermin.id] then unassign_unit_or_vermin(bld_assignments[unit_or_vermin.id], unit_or_vermin) end end function AssignAnimal:toggle_item_base(choice, target_value, bld_assignments) local true_value = self.status[self.status_revmap[1]].value if target_value == nil then target_value = choice.data.status ~= true_value end if target_value and choice.data.status == true_value then return target_value end if not target_value and choice.data.status ~= true_value then return target_value end if self.initial_min_disposition ~= DISPOSITION.PET.value and choice.data.disposition == DISPOSITION.PET.value then return target_value end local unit_or_vermin = choice.data.unit or choice.data.vermin detach_unit_or_vermin(unit_or_vermin, bld_assignments) if target_value then local displaced_unit = self.attach(unit_or_vermin) if displaced_unit then -- assigning a unit to a restraint can unassign a different unit for _, c in ipairs(self.subviews.list:getChoices()) do if c.data.unit == displaced_unit then c.data.status = self.get_status(displaced_unit) end end end end -- don't pass bld_assignments since it may no longer be valid choice.data.status = self.get_status(unit_or_vermin) return target_value end function AssignAnimal:select_item(idx, choice) if not dfhack.internal.getModifiers().shift then self.prev_list_idx = self.subviews.list.list:getSelected() end end function AssignAnimal:toggle_item(idx, choice) self:toggle_item_base(choice) end function AssignAnimal:toggle_range(idx, choice) if not self.get_multi_select() then self:toggle_item(idx, choice) return end if not self.prev_list_idx then self:toggle_item(idx, choice) return end local choices = self.subviews.list:getVisibleChoices() local list_idx = self.subviews.list.list:getSelected() local bld_assignments = get_bld_assignments() local target_value for i = list_idx, self.prev_list_idx, list_idx < self.prev_list_idx and 1 or -1 do target_value = self:toggle_item_base(choices[i], target_value, bld_assignments) end self.prev_list_idx = list_idx end function AssignAnimal:toggle_visible() local bld_assignments = get_bld_assignments() local target_value for _, choice in ipairs(self.subviews.list:getVisibleChoices()) do target_value = self:toggle_item_base(choice, target_value, bld_assignments) end end -- ------------------- -- AssignAnimalScreen -- view = view or nil AssignAnimalScreen = defclass(AssignAnimalScreen, gui.ZScreen) AssignAnimalScreen.ATTRS { focus_path='zone/assign', is_valid_ui_state=DEFAULT_NIL, status=DEFAULT_NIL, status_revmap=DEFAULT_NIL, get_status=DEFAULT_NIL, get_allow_vermin=DEFAULT_NIL, get_multi_select=DEFAULT_NIL, attach=DEFAULT_NIL, initial_min_disposition=DEFAULT_NIL, target_name=DEFAULT_NIL, } function AssignAnimalScreen:init() self:addviews{ AssignAnimal{ status=self.status, status_revmap=self.status_revmap, get_status=self.get_status, get_allow_vermin=self.get_allow_vermin, get_multi_select=self.get_multi_select, attach=self.attach, initial_min_disposition=self.initial_min_disposition, target_name=self.target_name, } } end function AssignAnimalScreen:onInput(keys) local handled = AssignAnimalScreen.super.onInput(self, keys) if not self.is_valid_ui_state() then if view then view:dismiss() end return end if keys._MOUSE_L then -- if any click is made outside of our window, we need to recheck unit properties local window = self.subviews[1] if not window:getMouseFramePos() then for _, choice in ipairs(self.subviews.list:getChoices()) do choice.data.status = self.get_status(choice.data.unit or choice.data.vermin) end window:refresh_list() end end return handled end function AssignAnimalScreen:onRenderFrame() if view and not self.is_valid_ui_state() then view:dismiss() end end function AssignAnimalScreen:onDismiss() view = nil end -- ------------------- -- PasturePondOverlay -- PasturePondOverlay = defclass(PasturePondOverlay, overlay.OverlayWidget) PasturePondOverlay.ATTRS{ desc='Adds a link to launch the animal assignment UI to pastures and ponds.', default_pos={x=7,y=13}, default_enabled=true, viewscreens={'dwarfmode/Zone/Some/Pen', 'dwarfmode/Zone/Some/Pond'}, frame={w=31, h=4}, } local function is_valid_zone() return df.global.game.main_interface.bottom_mode_selected == df.main_bottom_mode_type.ZONE and df.global.game.main_interface.civzone.cur_bld and (df.global.game.main_interface.civzone.cur_bld.type == df.civzone_type.Pen or df.global.game.main_interface.civzone.cur_bld.type == df.civzone_type.Pond) end local function is_pit_selected() return df.global.game.main_interface.bottom_mode_selected == df.main_bottom_mode_type.ZONE and df.global.game.main_interface.civzone.cur_bld and df.global.game.main_interface.civzone.cur_bld.type == df.civzone_type.Pond end local function attach_to_zone(unit_or_vermin) local zone = df.global.game.main_interface.civzone.cur_bld local ref = df.new(df.general_ref_building_civzone_assignedst) ref.building_id = zone.id; unit_or_vermin.general_refs:insert('#', ref) local is_unit = df.unit:is_instance(unit_or_vermin) local vec = is_unit and zone.assigned_units or zone.assigned_items utils.insert_sorted(vec, unit_or_vermin.id) end local PASTURE_STATUS = { NONE={label='Unknown', value=0}, ASSIGNED_HERE={label='Assigned here', value=1}, PASTURED={label='In other pasture', value=2}, PITTED={label='In other pit/pond', value=3}, RESTRAINED={label='On restraint', value=4}, BUILT_CAGE={label='In built cage', value=5}, ITEM_CAGE={label='In stockpiled cage', value=6}, ROAMING={label='Roaming', value=7}, } local PASTURE_STATUS_REVMAP = {} for k, v in pairs(PASTURE_STATUS) do PASTURE_STATUS_REVMAP[v.value] = k end local function get_zone_status(unit_or_vermin, bld_assignments) local assigned_zone_ref = get_general_ref(unit_or_vermin, df.general_ref_type.BUILDING_CIVZONE_ASSIGNED) if assigned_zone_ref then if df.global.game.main_interface.civzone.cur_bld.id == assigned_zone_ref.building_id then return PASTURE_STATUS.ASSIGNED_HERE.value else local civzone = df.building.find(assigned_zone_ref.building_id) if civzone.type == df.civzone_type.Pen then return PASTURE_STATUS.PASTURED.value elseif civzone.type == df.civzone_type.Pond then return PASTURE_STATUS.PITTED.value else return PASTURE_STATUS.NONE.value end end end if get_general_ref(unit_or_vermin, df.general_ref_type.BUILDING_CHAIN) then return PASTURE_STATUS.RESTRAINED.value end local cage_ref = get_cage_ref(unit_or_vermin) if cage_ref then if get_built_cage(df.item.find(cage_ref.item_id)) then return PASTURE_STATUS.BUILT_CAGE.value else return PASTURE_STATUS.ITEM_CAGE.value end end bld_assignments = bld_assignments or get_bld_assignments() if bld_assignments and bld_assignments[unit_or_vermin.id] then if df.building_chainst:is_instance(bld_assignments[unit_or_vermin.id]) then return PASTURE_STATUS.RESTRAINED.value end return PASTURE_STATUS.BUILT_CAGE.value end return PASTURE_STATUS.ROAMING.value end local function show_pasture_pond_screen() return AssignAnimalScreen{ is_valid_ui_state=is_valid_zone, status=PASTURE_STATUS, status_revmap=PASTURE_STATUS_REVMAP, get_status=get_zone_status, get_allow_vermin=is_pit_selected, get_multi_select=function() return true end, attach=attach_to_zone, target_name='pasture/pond/pit', }:show() end function PasturePondOverlay:init() self:addviews{ widgets.TextButton{ frame={t=0, l=0, w=23, h=1}, label='DFHack assign', key='CUSTOM_CTRL_T', on_activate=function() view = view and view:raise() or show_pasture_pond_screen() end, }, widgets.TextButton{ frame={b=0, l=0, w=28, h=1}, label='DFHack autobutcher', key='CUSTOM_CTRL_B', on_activate=function() dfhack.run_script('gui/autobutcher') end, }, } end -- ------------------- -- CageChainOverlay -- CageChainOverlay = defclass(CageChainOverlay, overlay.OverlayWidget) CageChainOverlay.ATTRS{ desc='Adds a link to launch the animal assignment UI to cages and chains.', default_pos={x=-40,y=34}, default_enabled=true, viewscreens={'dwarfmode/ViewSheets/BUILDING/Cage', 'dwarfmode/ViewSheets/BUILDING/Chain'}, frame={w=23, h=1}, frame_background=gui.CLEAR_PEN, } local function is_valid_building() local bld = dfhack.gui.getSelectedBuilding(true) if not bld or bld:getBuildStage() ~= bld:getMaxBuildStage() then return false end local bt = bld:getType() if bt ~= df.building_type.Cage and bt ~= df.building_type.Chain then return false end for _,zone in ipairs(bld.relations) do if zone.type == df.civzone_type.Dungeon then return false end end return true end local function is_cage_selected() local bld = dfhack.gui.getSelectedBuilding(true) return bld and bld:getType() == df.building_type.Cage end local function attach_to_building(unit_or_vermin) local bld = dfhack.gui.getSelectedBuilding(true) if not bld then return end local is_unit = df.unit:is_instance(unit_or_vermin) if is_unit and unit_or_vermin.relationship_ids[df.unit_relationship_type.Pet] ~= -1 then -- pet owners would just release them return end if bld:getType() == df.building_type.Cage then local vec = is_unit and bld.assigned_units or bld.assigned_items vec:insert('#', unit_or_vermin.id) elseif is_unit and bld:getType() == df.building_type.Chain then local prev_unit = bld.assigned bld.assigned = unit_or_vermin return prev_unit end end local function get_built_chain(unit_or_vermin) local built_chain_ref = get_general_ref(unit_or_vermin, df.general_ref_type.BUILDING_CHAIN) if not built_chain_ref then return end return df.building.find(built_chain_ref.building_id) end local CAGE_STATUS = { NONE={label='Unknown', value=0}, ASSIGNED_HERE={label='Assigned here', value=1}, PASTURED={label='In pasture', value=2}, PITTED={label='In pit/pond', value=3}, RESTRAINED={label='On other chain', value=4}, BUILT_CAGE={label='In built cage', value=5}, ITEM_CAGE={label='In stockpiled cage', value=6}, ROAMING={label='Roaming', value=7}, } local CAGE_STATUS_REVMAP = {} for k, v in pairs(CAGE_STATUS) do CAGE_STATUS_REVMAP[v.value] = k end local function get_cage_status(unit_or_vermin, bld_assignments) local bld = dfhack.gui.getSelectedBuilding(true) local is_chain = bld and df.building_chainst:is_instance(bld) if is_chain and bld.assigned == unit_or_vermin then return CAGE_STATUS.ASSIGNED_HERE.value elseif not is_chain and get_assigned_unit_or_vermin_idx(bld, unit_or_vermin) then return CAGE_STATUS.ASSIGNED_HERE.value end local cage_ref = get_cage_ref(unit_or_vermin) if cage_ref then if get_built_cage(df.item.find(cage_ref.item_id)) then return CAGE_STATUS.BUILT_CAGE.value end return CAGE_STATUS.ITEM_CAGE.value end local built_chain = get_built_chain(unit_or_vermin) if built_chain then if bld and bld == built_chain then return CAGE_STATUS.ASSIGNED_HERE.value end return CAGE_STATUS.RESTRAINED.value end local assigned_zone_ref = get_general_ref(unit_or_vermin, df.general_ref_type.BUILDING_CIVZONE_ASSIGNED) if assigned_zone_ref then local civzone = df.building.find(assigned_zone_ref.building_id) if civzone.type == df.civzone_type.Pen then return CAGE_STATUS.PASTURED.value elseif civzone.type == df.civzone_type.Pond then return CAGE_STATUS.PITTED.value else return CAGE_STATUS.NONE.value end end bld_assignments = bld_assignments or get_bld_assignments() if bld_assignments and bld_assignments[unit_or_vermin.id] then if df.building_chainst:is_instance(bld_assignments[unit_or_vermin.id]) then return CAGE_STATUS.RESTRAINED.value end return CAGE_STATUS.BUILT_CAGE.value end return CAGE_STATUS.ROAMING.value end local function show_cage_chain_screen() return AssignAnimalScreen{ is_valid_ui_state=is_valid_building, status=CAGE_STATUS, status_revmap=CAGE_STATUS_REVMAP, get_status=get_cage_status, get_allow_vermin=is_cage_selected, get_multi_select=is_cage_selected, attach=attach_to_building, initial_min_disposition=DISPOSITION.TAME.value, target_name='cage/restraint', }:show() end function CageChainOverlay:init() self:addviews{ widgets.TextButton{ frame={t=0, l=0, r=0, h=1}, label='DFHack assign', key='CUSTOM_CTRL_T', visible=is_valid_building, on_activate=function() view = view and view:raise() or show_cage_chain_screen() end, }, } end -- --------------------- -- RetireLocationOverlay -- local mi = df.global.game.main_interface local function location_details_is_on_top() return not mi.name_creator.open and not mi.unit_selector.open end RetireLocationOverlay = defclass(RetireLocationOverlay, overlay.OverlayWidget) RetireLocationOverlay.ATTRS{ desc='Adds a button to retire unneeded locations to the location details screen.', default_pos={x=-39,y=6}, default_enabled=true, viewscreens='dwarfmode/LocationDetails', frame={w=25, h=4}, } 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, subviews={ widgets.HotkeyLabel{ frame={t=0, l=0}, label='Retire location', key='CUSTOM_CTRL_D', on_activate=self:callback('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 details = mi.location_details local location = details.selected_ab local has_occupations, has_zones = false, #location.contents.building_ids ~= 0 for _, occupation in ipairs(location.occupations) do if occupation.histfig_id ~= -1 then has_occupations = true break end end if has_occupations or has_zones then local messages = {'Cannot retire location! Please:', ''} if has_occupations then table.insert(messages, '- unassign location occupations') end if has_zones then table.insert(messages, ('- detach this location from %d zone(s)'):format(#location.contents.building_ids)) end table.insert(messages, '') table.insert(messages, 'and try again') dialogs.showMessage('Location in use', table.concat(messages, NEWLINE)) return end 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) end OVERLAY_WIDGETS = { pasturepond=PasturePondOverlay, cagechain=CageChainOverlay, retirelocation=RetireLocationOverlay, } return _ENV