set-stress.lua

-- Sets stress to an arbitrary value
--By Putnam; http://www.bay12forums.com/smf/index.php?topic=139553.msg5820486#msg5820486
--Modified by /u/sesszet
--@module = true
local help = [====[

set-stress
=============
Sets stress to 0 or -value <value>; the normal range is 0 to 500,000 with very stable or
very stressed dwarves taking on negative or greater values respectively.
Applies to the selected unit, or use ``set-stress -all`` to apply to all units.
Use ``set-stress -max 20000`` to cap stress to 20000 (unhappy but bearable) for all units.

]====]

local utils = require 'utils'

function capStress(unit, value)
    if unit.counters.soldier_mood > df.unit.T_counters.T_soldier_mood.Enraged then
        -- Tantrum, Depressed, or Oblivious
        unit.counters.soldier_mood = df.unit.T_counters.T_soldier_mood.None
    end
    if unit.status.current_soul and unit.status.current_soul.personality.stress_level > tonumber(value) then
        unit.status.current_soul.personality.stress_level = value
    end
end

function setStress(unit, value)
    if unit.counters.soldier_mood > df.unit.T_counters.T_soldier_mood.Enraged then
        -- Tantrum, Depressed, or Oblivious
        unit.counters.soldier_mood = df.unit.T_counters.T_soldier_mood.None
    end
    if unit.status.current_soul then
        unit.status.current_soul.personality.stress_level = value
    end
end

local validArgs = utils.invert({
    'help',
        'max',
        'value',
    'all'
})

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

    if args.help then
        print(help)
        return
    end

    local value = args.value or 0
    local max = args.max or nil

    if args.all then
        for k,v in ipairs(df.global.world.units.active) do
                        if max then
                                capStress(v, max)
                        else
                                setStress(v, value)
                        end
        end
    else
        local unit = dfhack.gui.getSelectedUnit()
        if unit then
                        if max then
                                capStress(unit, max)
                        else
                                setStress(unit, value)
                        end
        else
            error 'Invalid usage: No unit selected and -all argument not given.'
        end
    end
end

if not dfhack_flags.module then
    main(...)
end