local _ENV = mkmodule('plugins.zone') 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 view:dismiss() return end if keys._MOUSE_L_DOWN 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 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{ 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{ 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) return bld and bld:getBuildStage() == bld:getMaxBuildStage() and (bld:getType() == df.building_type.Cage or bld:getType() == df.building_type.Chain) 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', on_activate=function() view = view and view:raise() or show_cage_chain_screen() end, }, } end OVERLAY_WIDGETS = { pasturepond=PasturePondOverlay, cagechain=CageChainOverlay, } return _ENV