-- Interface to assign caged units in a stockpile to pond/pit.

local utils = require 'utils'

local gui = require('gui')

local widgets = require('gui.widgets')

local function getCagedUnits(stockpile)

local cagedUnits = {}

for _, item in pairs(dfhack.buildings.getStockpileContents(stockpile)) do

if not dfhack.items.getGeneralRef(item, df.general_ref_type.CONTAINS_UNIT) then

goto skipitem

end

for _, ref in pairs(item.general_refs) do

if ref:getType() == df.general_ref_type.CONTAINS_UNIT then

table.insert(cagedUnits, ref.unit_id)

end

end

:: skipitem ::

end

return cagedUnits

end

Masspit = defclass(Masspit, widgets.Window)

Masspit.ATTRS {

frame_title='Mass-pit caged creatures',

frame={w=43, h=22, r = 2, t = 18},

resizable=true,

resize_min={h=12},

}

function Masspit:init()

self:addviews{

widgets.WrappedLabel{

view_id = 'label',

frame = { l = 0, t = 0 },

text_to_wrap = "Select or click on a stockpile"

},

widgets.Pages{

view_id ='pages',

subviews = {

widgets.List{

view_id = 'stockpiles',

frame = { l = 0, t = 4 },

on_select = self:callback("focusOnZone"),

on_submit = self:callback("setStockpile")

},

widgets.List{

view_id = 'pits',

frame = { l = 0, t = 3 },

on_select = self:callback("focusOnZone"),

on_submit = self:callback("setPit")

},

widgets.List{

view_id = 'pitted',

frame = { l = 0, t = 2 },

},

}

},

}

self:showStockpiles()

end

function Masspit:showStockpiles()

self.subviews.pages:setSelected('stockpiles')

local choices = {}

for _, zone in pairs(df.global.world.buildings.other.STOCKPILE) do

if (#zone.settings.animals.enabled > 0) then

if #getCagedUnits(zone) > 0 then

local zone_name = #zone.name > 0 and zone.name or ("Animal stockpile #" .. zone.stockpile_number)

local zone_position = df.coord:new()

zone_position.x = zone.centerx

zone_position.y = zone.centery

zone_position.z = zone.z

table.insert(choices, {

text = zone_name,

zone_position = zone_position,

zone_id = zone.id,

})

end

end

end

self.subviews.pages.subviews.stockpiles:setChoices(choices)

if #choices == 0 then

self.subviews.label.text_to_wrap = "No stockpiles with caged creatures found! Create one first."

else

self.subviews.label.text_to_wrap = "Select (or click on) an animal stockpile with caged creatures to pit."

end

end

local function getElevation()

return df.global.world.map.region_z + df.global.window_z - 100

end

function Masspit:showPondPits()

self.subviews.pages:setSelected('pits')

local choices = {}

for _, zone in pairs(df.global.world.buildings.other.ACTIVITY_ZONE) do

if (zone.type == df.civzone_type.Pond) then

local zone_name = #zone.name > 0 and zone.name or ("Unnamed pit/pond on elevation "..getElevation())

local zone_position = df.coord:new()

zone_position.x = zone.centerx

zone_position.y = zone.centery

zone_position.z = zone.z

table.insert(choices, {

text = zone_name,

zone_position = zone_position,

zone_id = zone.id

})

end

end

self.subviews.pages.subviews.pits:setChoices(choices)

if #choices == 0 then

self.subviews.label.text_to_wrap = "No pond/pit zones found! Create one first."

else

self.subviews.label.text_to_wrap = ([[Select (or click on) a pond/pit zone to move %s caged creatures to.]]):format(#self.caged_units)

end

self:updateLayout()

end

function Masspit:focusOnZone(index, choice)

if not index or not choice then

return

end

if (choice.zone_position) then

dfhack.gui.revealInDwarfmodeMap(choice.zone_position)

end

end

function Masspit:setStockpile(_, choice)

self.stockpile = utils.binsearch(df.global.world.buildings.other.STOCKPILE, choice.zone_id, 'id')

self.caged_units = getCagedUnits(self.stockpile)

if #self.caged_units == 0 then

self.subviews.label.text_to_wrap = "This stockpile contains no caged units."

self.subviews.label:updateLayout()

else

self:showPondPits()

end

end

function Masspit:setPit(_, choice)

self.pit = utils.binsearch(df.global.world.buildings.other.ACTIVITY_ZONE, choice.zone_id, 'id')

self.subviews.label.text_to_wrap = ([[Pitting %s units.]]):format(#self.caged_units)

self.subviews.label:updateLayout()

local choices = {}

for _, unit_id in pairs(self.caged_units) do

local unit = df.unit.find(unit_id)

local unit_name = dfhack.units.getReadableName(unit)

-- Prevents duplicate units in assignments, which can cause crashes.

local duplicate = false

for _, assigned_unit in pairs(self.pit.assigned_units) do

if assigned_unit == unit.id then

duplicate = true

end

end

table.insert(choices, { text = ([[%s: %s%s]]):format(unit.id, unit_name, duplicate and " (ALREADY ASSIGNED)" or "") })

if not duplicate then

unit.general_refs:insert("#", { new = df.general_ref_building_civzone_assignedst, building_id=self.pit.id })

self.pit.assigned_units:insert("#", unit.id)

end

end

self.subviews.pages.subviews.pitted:setChoices(choices)

self.subviews.pages:setSelected('pitted')

end

function Masspit:onInput(keys)

if Masspit.super.onInput(self, keys) then return true end

if keys._MOUSE_L and not self:getMouseFramePos() then

if self.subviews.pages:getSelected() == 1 then

local building = dfhack.buildings.findAtTile(dfhack.gui.getMousePos())

if building and df.building_stockpilest:is_instance(building) then

self:setStockpile(nil, { zone_id = building.id })

end

elseif self.subviews.pages:getSelected() == 2 then

local zones = dfhack.buildings.findCivzonesAt(dfhack.gui.getMousePos())

for _, zone in pairs(zones) do

if (zone.type == df.civzone_type.Pond) then

self:setPit(nil, { zone_id = zone.id })

end

end

end

end

end

-- Screen setup

MasspitScreen = defclass(MasspitScreen, gui.ZScreen)

MasspitScreen.ATTRS {

focus_path='masspit',

pass_movement_keys=true,

}

function MasspitScreen:init()

self:addviews{Masspit{}}

end

function MasspitScreen:onDismiss()

view = nil

end

if not dfhack.world.isFortressMode() or not dfhack.isMapLoaded() then

qerror('This script requires a fortress map to be loaded')

end

view = view and view:raise() or MasspitScreen{}:show()