|
| 1 | +-- Consolidate and remove extra ammo items to fix 'Soldier (no item)' issue. |
| 2 | + |
| 3 | +local argparse = require("argparse") |
| 4 | +local utils = require('utils') |
| 5 | + |
| 6 | +local function GetTrainingSquads() |
| 7 | + local trainingSquads = {} |
| 8 | + for _, squad in ipairs(df.global.world.squads.all) do |
| 9 | + if squad.entity_id == df.global.plotinfo.group_id then |
| 10 | + if #squad.ammo.ammunition > 0 and squad.activity ~= -1 then |
| 11 | + trainingSquads[#trainingSquads + 1] = squad |
| 12 | + end |
| 13 | + end |
| 14 | + end |
| 15 | + return trainingSquads |
| 16 | +end |
| 17 | + |
| 18 | +local function isTrainingAmmo(ammoItem, squad) |
| 19 | + for _, ammoSpec in ipairs(squad.ammo.ammunition) do |
| 20 | + if ammoSpec.flags.use_training then |
| 21 | + for _, id in ipairs(ammoSpec.assigned) do |
| 22 | + if ammoItem.id == id then return true end |
| 23 | + end |
| 24 | + end |
| 25 | + end |
| 26 | + return false |
| 27 | +end |
| 28 | + |
| 29 | +local function GetTrainingAmmo(quiver, squad) |
| 30 | + local trainingAmmo = {} |
| 31 | + for _, generalRef in ipairs(quiver.general_refs) do |
| 32 | + if df.general_ref_contains_itemst:is_instance(generalRef) then |
| 33 | + local containedAmmo = generalRef |
| 34 | + local ammoItem = containedAmmo and df.item.find(containedAmmo.item_id) |
| 35 | + if isTrainingAmmo(ammoItem, squad) then |
| 36 | + trainingAmmo[#trainingAmmo + 1] = ammoItem |
| 37 | + end |
| 38 | + end |
| 39 | + end |
| 40 | + return trainingAmmo |
| 41 | +end |
| 42 | + |
| 43 | +local function UnassignAmmo(trainingAmmo, itemToKeep, itemsToRemove, squad, unit) |
| 44 | + local plotEqAssignedAmmo = df.global.plotinfo.equipment.items_assigned.AMMO |
| 45 | + local plotEqUnassignedAmmo = df.global.plotinfo.equipment.items_unassigned.AMMO |
| 46 | + local uniforms = { |
| 47 | + unit.uniform.uniforms.CLOTHING, |
| 48 | + unit.uniform.uniforms.REGULAR, |
| 49 | + unit.uniform.uniforms.TRAINING, |
| 50 | + unit.uniform.uniforms.TRAINING_RANGED |
| 51 | + } |
| 52 | + for _, ammoItem in ipairs(trainingAmmo) do |
| 53 | + if ammoItem ~= itemToKeep then |
| 54 | + local idx |
| 55 | + local assignedAmmo |
| 56 | + for _, ammoSpec in ipairs(squad.ammo.ammunition) do |
| 57 | + if ammoSpec.flags.use_training then |
| 58 | + idx = utils.linear_index(ammoSpec.assigned, ammoItem.id) |
| 59 | + if idx then |
| 60 | + assignedAmmo = ammoSpec.assigned |
| 61 | + goto unassignAmmo |
| 62 | + end |
| 63 | + end |
| 64 | + end |
| 65 | + ::unassignAmmo:: |
| 66 | + if assignedAmmo and idx then |
| 67 | + -- Unassign ammo item from squad. |
| 68 | + assignedAmmo:erase(idx) |
| 69 | + idx = utils.linear_index(squad.ammo.ammo_items, ammoItem.id) |
| 70 | + if idx then |
| 71 | + -- Remove item/unit pairings. |
| 72 | + squad.ammo.ammo_items:erase(idx) |
| 73 | + squad.ammo.ammo_units:erase(idx) |
| 74 | + end |
| 75 | + idx = utils.linear_index(plotEqAssignedAmmo, ammoItem.id) |
| 76 | + if idx then |
| 77 | + -- Move ammo item from assigned ammo list to unassigned ammo list. |
| 78 | + plotEqAssignedAmmo:erase(idx) |
| 79 | + plotEqUnassignedAmmo:insert('#', ammoItem.id) |
| 80 | + utils.sort_vector(plotEqUnassignedAmmo) |
| 81 | + end |
| 82 | + end |
| 83 | + for _, uniform in ipairs(uniforms) do |
| 84 | + -- Remove ammo item from uniform. |
| 85 | + idx = utils.linear_index(uniform, ammoItem.id) |
| 86 | + if idx then uniform:erase(idx) end |
| 87 | + end |
| 88 | + if not utils.linear_index(itemsToRemove, ammoItem) then |
| 89 | + -- Force drop ammo item to avoid issue recurring if game reassigns the ammo item to squad. |
| 90 | + -- unit.uniform.uniform_drop:insert('#', ammoItem.id) |
| 91 | + -- Units that choose to haul the surplus ammo items to stockpiles instead of just dropping them |
| 92 | + -- on the ground will cancel their archery practice and put away the ammo item they were supposed |
| 93 | + -- to train with as well. Force dropping the surplus item with moveToGround circumvents this. |
| 94 | + local pos = unit and xyz2pos(dfhack.units.getPosition(unit)) |
| 95 | + dfhack.items.moveToGround(ammoItem, pos) |
| 96 | + end |
| 97 | + end |
| 98 | + end |
| 99 | + -- Prompt unit to drop item. |
| 100 | + -- unit.uniform.pickup_flags.update = true |
| 101 | +end |
| 102 | + |
| 103 | +-- For practicality, item material, quality, and its creator (for masterworks), is ignored |
| 104 | +-- for the purpose of combining the limited number of ammo items inside a quiver. |
| 105 | +local function ConsolidateAmmo(trainingAmmo, squad, unit) |
| 106 | + local itemToKeep |
| 107 | + local itemsToRemove = {} |
| 108 | + -- Check first if any training ammo item already has a stack size of 25 or higher. |
| 109 | + for _, ammoItem in ipairs(trainingAmmo) do |
| 110 | + if ammoItem.stack_size >= 25 then |
| 111 | + itemToKeep = ammoItem |
| 112 | + goto unassignAmmo |
| 113 | + end |
| 114 | + end |
| 115 | + for _, ammoItem in ipairs(trainingAmmo) do |
| 116 | + if not itemToKeep then |
| 117 | + -- Keep the first item. |
| 118 | + itemToKeep = ammoItem |
| 119 | + goto nextItem |
| 120 | + end |
| 121 | + if itemToKeep and ammoItem ~= itemToKeep and itemToKeep.stack_size < 25 then |
| 122 | + local combineSize = 25 - itemToKeep.stack_size |
| 123 | + if ammoItem.stack_size > combineSize then |
| 124 | + itemToKeep.stack_size = itemToKeep.stack_size + combineSize |
| 125 | + ammoItem.stack_size = ammoItem.stack_size - combineSize |
| 126 | + else |
| 127 | + itemToKeep.stack_size = itemToKeep.stack_size + ammoItem.stack_size |
| 128 | + itemsToRemove[#itemsToRemove + 1] = ammoItem |
| 129 | + end |
| 130 | + end |
| 131 | + ::nextItem:: |
| 132 | + end |
| 133 | + ::unassignAmmo:: |
| 134 | + -- Unassign surplus ammo items first before removing any from the game. |
| 135 | + UnassignAmmo(trainingAmmo, itemToKeep, itemsToRemove, squad, unit) |
| 136 | + if #itemsToRemove > 0 then |
| 137 | + for _, item in ipairs(itemsToRemove) do |
| 138 | + dfhack.items.remove(item) |
| 139 | + end |
| 140 | + end |
| 141 | +end |
| 142 | + |
| 143 | +local function FixTrainingUnits(trainingSquads, options) |
| 144 | + local totalTrainingAmmo = 0 |
| 145 | + local consolidateCount = 0 |
| 146 | + for _, squad in ipairs(trainingSquads) do |
| 147 | + for _, position in ipairs(squad.positions) do |
| 148 | + if position.occupant == -1 then goto nextPosition end |
| 149 | + local unit = df.unit.find(df.historical_figure.find(position.occupant).unit_id) |
| 150 | + local quiver = unit and df.item.find(position.equipment.quiver) |
| 151 | + if quiver then |
| 152 | + local trainingAmmo = GetTrainingAmmo(quiver, squad) |
| 153 | + if #trainingAmmo > 1 then |
| 154 | + if not options.quiet then |
| 155 | + local unitName = unit and dfhack.units.getReadableName(unit) |
| 156 | + print(('Consolidating training ammo for %s...'):format(unitName)) |
| 157 | + end |
| 158 | + totalTrainingAmmo = totalTrainingAmmo + #trainingAmmo |
| 159 | + ConsolidateAmmo(trainingAmmo, squad, unit) |
| 160 | + consolidateCount = consolidateCount + 1 |
| 161 | + end |
| 162 | + end |
| 163 | + ::nextPosition:: |
| 164 | + end |
| 165 | + end |
| 166 | + if not options.quiet then |
| 167 | + if consolidateCount > 0 then |
| 168 | + print(('%d stacks of ammo items in %d quiver(s) consolidated.'):format(totalTrainingAmmo, consolidateCount)) |
| 169 | + else |
| 170 | + print('No stacks of ammo items require consolidation.') |
| 171 | + end |
| 172 | + end |
| 173 | +end |
| 174 | + |
| 175 | +local function ParseCommandLine(args) |
| 176 | + local options = { |
| 177 | + help = false, |
| 178 | + quiet = false |
| 179 | + } |
| 180 | + local positionals = argparse.processArgsGetopt(args, { |
| 181 | + {'h', 'help', handler = function() options.help = true end}, |
| 182 | + {'q', 'quiet', handler=function() options.quiet = true end} |
| 183 | + }) |
| 184 | + return options |
| 185 | +end |
| 186 | + |
| 187 | +local function Main(args) |
| 188 | + local options = ParseCommandLine(args) |
| 189 | + if args[1] == 'help' or options.help then |
| 190 | + print(dfhack.script_help()) |
| 191 | + return |
| 192 | + end |
| 193 | + local trainingSquads = GetTrainingSquads() |
| 194 | + if #trainingSquads < 1 then |
| 195 | + if not options.quiet then print('No ranged squads are currently training.') end |
| 196 | + return |
| 197 | + end |
| 198 | + FixTrainingUnits(trainingSquads, options) |
| 199 | +end |
| 200 | + |
| 201 | +if not dfhack.isSiteLoaded() and not dfhack.world.isFortressMode() then |
| 202 | + qerror('This script requires the game to be in fortress mode.') |
| 203 | +end |
| 204 | + |
| 205 | +Main({...}) |
0 commit comments