dfhack/plugins/lua/zone.lua

1178 lines
42 KiB
Lua

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