User:Fleeting Frames/militarygearclaims

Save as militarygearclaims.lua into /hack/scripts/fix folder.

Finds units who have claimed nonspecific gear that is inaccessible/forbidden. Then removes them; resulting in military picking up new items once month changes. Exludes items in inventories unless otherwise specified. Dwarves seem to go back to same cloth items, though this does not prevent them from wearing other ones, as it does with military gear.

Options:

-info: only prints info about claimed gear and does nothing

-textless: only removes unusable gear assignments and prints nothing

-textlessifempty: only prints if it changed something

-translatename: prints translated names instead

-removefrominv also removes {items} in inventory

-ignoresquad: ignores items assigned as part of uniform

-ignoreunit: ignores clothing not part of the uniform

-selected: apply only to currently selected unit

-help: display this text

-- Unassigns forbidden/inaccessible claimed military gear
local help = [====[

fix/militarygearclaims
======================
Finds units who have claimed nonspecific gear that is inaccessible/forbidden.
Then removes them; resulting in military picking up new items once month changes.
Exludes carried items. Dwarves seem to go back to same cloth items, though this
does not prevent them from wearing other ones, as it does with military gear.

Options:
:-info:            only prints info about claimed gear and does nothing
:-textless:        only removes unusable gear assignments and prints nothing
:-textlessifempty: only prints if it changed something
:-translatename:   prints translated names instead
:-removefrominv    also removes {items} in inventory
:-ignoresquad:     ignores items assigned as part of uniform
:-ignoreunit:      ignores clothing not part of the uniform
:-selected:        apply only to currently selected unit
:-help:            display this text
]====]

local utils = require('utils')

validArgs = utils.invert({
  'info',
  'textless',
  'textlessifempty',
  'translatename',
  'removefrominv',
  'ignoresquad',
  'ignoreunit',
  'selected',
  'help'
})

local args = utils.processArgs({...}, validArgs)

if args.help then print(help) return end

local uniformcategoryfromtype = {
    [df.item_armorst] = "body",
    [df.item_helmst] = "head",
    [df.item_pantsst] = "pants",
    [df.item_glovesst] = "gloves",
    [df.item_shoesst] = "shoes",
    [df.item_shieldst] = "shield",
    [df.item_weaponst] = "weapon"
}

local misccategoryfromtype = {
    [df.item_quiverst] = "quiver",
    [df.item_flaskst] = "flask",
    [df.item_backpackst] = "backpack"

}


function getFlag(dfvalue, key)
    -- Utility function for safely requesting info from df data
    if not dfvalue or not key then return nil end --pairs crash prevention
    flagtypetable = flagtypetable or {} --memoization so it doesn't iterate through the dfvalue if we've already checked that type of value for given flag
    if flagtypetable[dfvalue._type] and flagtypetable[dfvalue._type][key] then return dfvalue[key] end
    if flagtypetable[dfvalue._type] and flagtypetable[dfvalue._type][key]==false then return nil end
    if not flagtypetable[dfvalue._type] then flagtypetable[dfvalue._type] = {} end
    for akey, avalue in pairs(dfvalue) do
        if akey == key then
            flagtypetable[dfvalue._type][key] = true
            return dfvalue[akey] 
        end
    end
    flagtypetable[dfvalue._type][key] = false
end


function getItemName(item)
    local retval = dfhack.items.getDescription(item,3)
    qualitytable = qualitytable or {
    [0] = " ",
    [1] = "-",
    [2] = "+",
    [3] = "\42",
    [4] = "\240",
    [5] = "\15"
    }
    function wrap(str,brackets)
        return brackets .. str .. brackets
    end
    retval = wrap(retval,qualitytable[item.quality])
    if #item.improvements > 0 then
        retval = "\174" .. retval .. "\175"
        local bestdecoration = 0
        for index, decoration in pairs(item.improvements) do
            if decoration.quality > bestdecoration then bestdecoration = decoration.quality end
        end
        retval = wrap(retval,qualitytable[bestdecoration])
        if item.flags.foreign then retval = "("..retval..")" end
    else
        if item.flags.foreign then retval = "(  "..retval.."  )"
        else retval = wrap(retval, "   ") end
    end
    return retval
end

local function getForbiddenUnitGear(unit, doErase)
    local ret, doreturn = {}
    for cur_uniform=0, #unit.military.uniforms-1 do --unit.military.cur_uniform do
        for i=#unit.military.uniforms[cur_uniform]-1, 0, -1 do 
            local item=df.item.find(unit.military.uniforms[cur_uniform][i])
            if item and 
               (
               ((not item.flags.in_inventory) and 
               (item.flags.forbid or not dfhack.maps.canWalkBetween(unit.pos,item.pos)))
               or (args.removefrominv and
               item.flags.in_inventory and
               item.flags.forbid)
               )
               then
                if not ret[item.id] then
                    local reti={}
                    reti.itemname = getItemName(item)
                    local pos = item.flags.in_inventory and xyz2pos(dfhack.items.getPosition(item)) or item.pos
                    reti.posstring = "x " .. pos.x .. " y " .. pos.y .. " z " .. pos.z
                    reti.id=item.id
                    reti.unitname = dfhack.df2console(dfhack.TranslateName(unit.name, args.translatename and true))
                    reti.type = uniformcategoryfromtype[item._type]
                    reti.generaltype = misccategoryfromtype[item._type]
                    ret[item.id]=reti
                    doreturn = true
                end
                if doErase then
                    unit.military.uniforms[cur_uniform]:erase(i)
                    for j=#unit.military.uniform_pickup-1,0,-1 do
                        if unit.military.uniform_pickup[j] == item.id then
                            unit.military.uniform_pickup:erase(j)
                        end
                    end
                end

            end
        end
    end
    return doreturn and ret or nil
end


local function isAssignableCitizen(unit)
    if dfhack.units.isOwnCiv(unit) and 
        -- dfhack.units.isOwnGroup(unit) and --prevents mercs
        (not
       ((getFlag(unit.flags1,'dead') or getFlag(unit.flags1,'inactive'))or unit.flags2.visitor or unit.flags3.ghostly) ) and
       df.profession.attrs[unit.profession].can_assign_labor 
       then
        return true
    end
end

local function getAllForbiddenGear(eraseInUnit, eraseInSquad)
    local lst = {}
    for index, unit in pairs(df.global.world.units.all) do
        if args.selected then unit = dfhack.gui.getSelectedUnit() end
        if isAssignableCitizen(unit) then
            local forbiddenUnitGear = getForbiddenUnitGear(unit,eraseInUnit)
            if forbiddenUnitGear then
                for _, forbiddenItem in pairs(forbiddenUnitGear) do
                    table.insert(lst, forbiddenItem)
                    if eraseInSquad and unit.military.squad_id > -1 then
                        local squadpos = df.squad.find(unit.military.squad_id).positions[unit.military.squad_position]
                        local specificitem = false
                        if forbiddenItem.type and #squadpos.uniform[forbiddenItem.type] > 0 then
                            for index, spec in pairs(squadpos.uniform[forbiddenItem.type]) do
                                if #spec.assigned>0 and spec.assigned[0] == forbiddenItem.id then 
                                    if spec.item == -1 then
                                        spec.assigned:erase(0)
                                    else
                                        specificitem = true
                                    end
                                end
                            end
                        elseif forbiddenItem.generaltype and squadpos[forbiddenItem.generaltype] > -1 then
                            squadpos[forbiddenItem.generaltype] = -1
                        end
                        
                        if not specificitem then
                            for i=#squadpos.assigned_items-1, 0, -1 do
                                if squadpos.assigned_items[i] == forbiddenItem.id then
                                    squadpos.assigned_items:erase(i)
                                end
                            end
                        else
                            forbiddenItem.itemname = forbiddenItem.itemname .. " in specific"
                        end
                    end
                end
            end
        end
        if args.selected then break end
    end
    return lst
end

if not args.textless and not args.textlessifempty then print("Looking for forbidden/inaccessible claimed gear...") end

local forbidgear
if args.info then
    forbidgear = getAllForbiddenGear(false,false)
else
    forbidgear = getAllForbiddenGear(not args.ignoreunit and true or false, not args.ignoresquad and true or false)
end

function pad(text, length)
    if #text < length then
        return pad(text .. " ", length)
    end
    return text, length
end

if not args.textless then
    if not (0==#forbidgear and args.textlessifempty) then
    if args.textlessifempty then print("Looking for forbidden/inaccessible claimed gear...") end
    
    for _, item in pairs(forbidgear) do
        print(dfhack.df2console(pad(item.itemname,42)), "id " .. item.id, item.posstring, " by " .. item.unitname)
    end
    
    if 0==#forbidgear then print("Found no forbidden/inaccessible claimed gear!") else print("Finished!") end
    end
end