|
| 1 | +--@module = true |
| 2 | + |
| 3 | +local argparse = require('argparse') |
| 4 | +local exterminate = reqscript('exterminate') |
| 5 | + |
| 6 | +local GLOBAL_KEY = 'fix/wildlife' |
| 7 | + |
| 8 | +DEBUG = DEBUG or false |
| 9 | + |
| 10 | +stuck_creatures = stuck_creatures or {} |
| 11 | + |
| 12 | +dfhack.onStateChange[GLOBAL_KEY] = function(sc) |
| 13 | + if (sc == SC_MAP_UNLOADED or sc == SC_MAP_LOADED) and |
| 14 | + dfhack.world.isFortressMode() |
| 15 | + then |
| 16 | + stuck_creatures = {} |
| 17 | + end |
| 18 | +end |
| 19 | + |
| 20 | +local function print_summary(opts, unstuck) |
| 21 | + if not next(unstuck) then |
| 22 | + if not opts.quiet then |
| 23 | + print('No stuck wildlife found') |
| 24 | + return |
| 25 | + end |
| 26 | + end |
| 27 | + local prefix = opts.week and (GLOBAL_KEY .. ': ') or '' |
| 28 | + local msg_txt = opts.dry_run and '' or 'no longer ' |
| 29 | + for _,entry in pairs(unstuck) do |
| 30 | + if entry.count == 1 then |
| 31 | + print(('%s%d %s is %sblocking new waves of wildlife'):format( |
| 32 | + prefix, |
| 33 | + entry.count, |
| 34 | + entry.known and dfhack.units.getRaceReadableNameById(entry.race) or 'hidden creature', |
| 35 | + msg_txt)) |
| 36 | + else |
| 37 | + print(('%s%d %s are %sblocking new waves of wildlife'):format( |
| 38 | + prefix, |
| 39 | + entry.count, |
| 40 | + entry.known and dfhack.units.getRaceNamePluralById(entry.race) or 'hidden creatures', |
| 41 | + msg_txt)) |
| 42 | + end |
| 43 | + end |
| 44 | +end |
| 45 | + |
| 46 | +local function refund_population(entry) |
| 47 | + local epop = entry.pop |
| 48 | + for _,population in ipairs(df.global.world.populations) do |
| 49 | + local wpop = population.population |
| 50 | + if population.quantity < 10000001 and |
| 51 | + wpop.region_x == epop.region_x and |
| 52 | + wpop.region_y == epop.region_y and |
| 53 | + wpop.feature_idx == epop.feature_idx and |
| 54 | + wpop.cave_id == epop.cave_id and |
| 55 | + wpop.site_id == epop.site_id and |
| 56 | + wpop.population_idx == epop.population_idx |
| 57 | + then |
| 58 | + population.quantity = math.min(population.quantity + entry.count, population.quantity_max) |
| 59 | + break |
| 60 | + end |
| 61 | + end |
| 62 | +end |
| 63 | + |
| 64 | +-- refund unit to population and ensure it doesn't get picked up by unstick_surface_wildlife in the future |
| 65 | +local function detach_unit(unit) |
| 66 | + unit.flags2.roaming_wilderness_population_source = false |
| 67 | + unit.flags2.roaming_wilderness_population_source_not_a_map_feature = false |
| 68 | + refund_population{race=unit.race, pop=unit.animal.population, known=true, count=1} |
| 69 | +end |
| 70 | + |
| 71 | +local TICKS_PER_DAY = 1200 |
| 72 | +local TICKS_PER_WEEK = TICKS_PER_DAY * 7 |
| 73 | +local TICKS_PER_MONTH = 28 * TICKS_PER_DAY |
| 74 | +local TICKS_PER_SEASON = 3 * TICKS_PER_MONTH |
| 75 | +local TICKS_PER_YEAR = 4 * TICKS_PER_SEASON |
| 76 | + |
| 77 | +local WEEK_BEFORE_EOY_TICKS = TICKS_PER_YEAR - TICKS_PER_WEEK |
| 78 | + |
| 79 | +-- update stuck_creatures records and check timeout |
| 80 | +-- we only enter this function if the unit's leave_countdown has already expired |
| 81 | +-- returns true if the unit has timed out |
| 82 | +local function check_timeout(opts, unit, week_ago_ticks) |
| 83 | + if not opts.week then return true end |
| 84 | + if not stuck_creatures[unit.id] then |
| 85 | + stuck_creatures[unit.id] = df.global.cur_year_tick |
| 86 | + return false |
| 87 | + end |
| 88 | + local timestamp = stuck_creatures[unit.id] |
| 89 | + return timestamp < week_ago_ticks or |
| 90 | + (timestamp > df.global.cur_year_tick and timestamp > WEEK_BEFORE_EOY_TICKS) |
| 91 | +end |
| 92 | + |
| 93 | +local function to_key(pop) |
| 94 | + return ('%d:%d:%d:%d:%d:%d'):format( |
| 95 | + pop.region_x, pop.region_y, pop.feature_idx, pop.cave_id, pop.site_id, pop.population_idx) |
| 96 | +end |
| 97 | + |
| 98 | +local function is_active_wildlife(unit) |
| 99 | + return not dfhack.units.isDead(unit) and |
| 100 | + dfhack.units.isActive(unit) and |
| 101 | + dfhack.units.isWildlife(unit) and |
| 102 | + unit.flags2.roaming_wilderness_population_source |
| 103 | +end |
| 104 | + |
| 105 | +-- called by force for the "Wildlife" event |
| 106 | +function free_all_wildlife(include_hidden) |
| 107 | + for _,unit in ipairs(df.global.world.units.active) do |
| 108 | + if is_active_wildlife(unit) and |
| 109 | + (include_hidden or not dfhack.units.isHidden(unit)) |
| 110 | + then |
| 111 | + detach_unit(unit) |
| 112 | + end |
| 113 | + end |
| 114 | +end |
| 115 | + |
| 116 | +local function unstick_surface_wildlife(opts) |
| 117 | + local unstuck = {} |
| 118 | + local week_ago_ticks = math.max(0, df.global.cur_year_tick - TICKS_PER_WEEK) |
| 119 | + for _,unit in ipairs(df.global.world.units.active) do |
| 120 | + if not is_active_wildlife(unit) or unit.animal.leave_countdown > 0 then |
| 121 | + goto skip |
| 122 | + end |
| 123 | + if not check_timeout(opts, unit, week_ago_ticks) then |
| 124 | + goto skip |
| 125 | + end |
| 126 | + local pop = unit.animal.population |
| 127 | + local unstuck_entry = ensure_key(unstuck, to_key(pop), {race=unit.race, pop=pop, known=false, count=0}) |
| 128 | + unstuck_entry.known = unstuck_entry.known or not dfhack.units.isHidden(unit) |
| 129 | + unstuck_entry.count = unstuck_entry.count + 1 |
| 130 | + if not opts.dry_run then |
| 131 | + stuck_creatures[unit.id] = nil |
| 132 | + exterminate.killUnit(unit, exterminate.killMethod.DISINTEGRATE) |
| 133 | + end |
| 134 | + ::skip:: |
| 135 | + end |
| 136 | + for _,entry in pairs(unstuck) do |
| 137 | + refund_population(entry) |
| 138 | + end |
| 139 | + print_summary(opts, unstuck) |
| 140 | +end |
| 141 | + |
| 142 | +if dfhack_flags.module then |
| 143 | + return |
| 144 | +end |
| 145 | + |
| 146 | +if not dfhack.world.isFortressMode() or not dfhack.isMapLoaded() then |
| 147 | + qerror('needs a loaded fortress map to work') |
| 148 | +end |
| 149 | + |
| 150 | +local opts = { |
| 151 | + dry_run=false, |
| 152 | + help=false, |
| 153 | + quiet=false, |
| 154 | + week=false, |
| 155 | +} |
| 156 | + |
| 157 | +local positionals = argparse.processArgsGetopt({...}, { |
| 158 | + {'h', 'help', handler = function() opts.help = true end}, |
| 159 | + {'n', 'dry-run', handler = function() opts.dry_run = true end}, |
| 160 | + {'w', 'week', handler = function() opts.week = true end}, |
| 161 | + {'q', 'quiet', handler = function() opts.quiet = true end}, |
| 162 | +}) |
| 163 | + |
| 164 | +if positionals[1] == 'help' or opts.help then |
| 165 | + print(dfhack.script_help()) |
| 166 | + return |
| 167 | +end |
| 168 | + |
| 169 | +if positionals[1] == 'ignore' then |
| 170 | + local unit |
| 171 | + local unit_id = positionals[2] and argparse.nonnegativeInt(positionals[2], 'unit_id') |
| 172 | + if unit_id then |
| 173 | + unit = df.unit.find(unit_id) |
| 174 | + else |
| 175 | + unit = dfhack.gui.getSelectedUnit(true) |
| 176 | + end |
| 177 | + if not unit then |
| 178 | + qerror('please select a unit or pass a unit ID on the commandline') |
| 179 | + end |
| 180 | + if not is_active_wildlife(unit) then |
| 181 | + qerror('selected unit is not blocking new waves of wildlife; nothing to do') |
| 182 | + end |
| 183 | + detach_unit(unit) |
| 184 | + if not opts.quiet then |
| 185 | + print(('%s will now be ignored by fix/wildlife'):format(dfhack.units.getReadableName(unit))) |
| 186 | + end |
| 187 | +else |
| 188 | + unstick_surface_wildlife(opts) |
| 189 | +end |
0 commit comments