full-heal.lua

--@ module = true

local utils = require('utils')

local validArgs = utils.invert({
    'r',
    'help',
    'unit',
    'keep_corpse',
    'all',
    'all_civ',
    'all_citizens'
})

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

if args.help then
    print(dfhack.script_help())
    return
end

function addResurrectionEvent(histFigID)
    local event = df.history_event_hist_figure_revivedst:new()
    event.histfig = histFigID
    event.year = df.global.cur_year
    event.seconds = df.global.cur_year_tick
    event.id = df.global.hist_event_next_id
    df.global.world.history.events:insert('#', event)
    df.global.hist_event_next_id = df.global.hist_event_next_id + 1
end

function heal(unit, resurrect, keep_corpse)
    if not unit then
        return
    end
    if resurrect then
        if unit.flags2.killed and not unit.flags3.scuttle then -- scuttle appears to be applicable to just wagons, which probably shouldn't be resurrected
            print(('Resurrecting %s'):format(dfhack.units.getReadableName(unit)))
            unit.flags1.inactive = false
            unit.flags2.slaughter = false
            unit.flags2.killed = false
            unit.flags3.ghostly = false -- TO DO: check whether ghost is currently walking through walls before removing ghostliness
            unit.ghost_info = nil

            local _, occupancy = dfhack.maps.getTileFlags(pos2xyz(unit.pos))
            occupancy.unit_grounded = true -- it's simpler to always set the tile occupancy to unit_grounded, as determining whether occupancy.unit is correct would require one to check for the presence of other units on that tile
            unit.flags1.on_ground = true -- this also fits the script thematically

            if unit.hist_figure_id ~= -1 then
                local hf = df.historical_figure.find(unit.hist_figure_id)
                hf.died_year = -1
                hf.died_seconds = -1
                hf.flags.ghost = false
                addResurrectionEvent(hf.id)
            end

            if dfhack.world.isFortressMode() and dfhack.units.isOwnCiv(unit) then
                unit.flags2.resident = false -- appears to be set to true for dead citizens in a reclaimed fortress, which causes them to be marked as hostile when resurrected

                local deadCitizens = df.global.plotinfo.main.dead_citizens
                for i = #deadCitizens-1,0,-1 do
                    if deadCitizens[i].unit_id == unit.id then
                        deadCitizens:erase(i)
                    end
                end
            end

            if not keep_corpse then
                local corpses = df.global.world.items.other.CORPSE
                for i = #corpses-1,0,-1 do
                    local corpse = corpses[i] --as:df.item_body_component
                    if corpse.unit_id == unit.id then
                        dfhack.items.remove(corpse)
                    end
                end
            end
        end
    end

    --print("Erasing wounds...")
    for _, wound in ipairs(unit.body.wounds) do
        wound:delete()
    end
    unit.body.wounds:resize(0)
    unit.body.wound_next_id = 1

    --print("Refilling blood...")
    unit.body.blood_count = unit.body.blood_max

    --print("Resetting grasp/stand/fly status...")
    unit.status2.limbs_stand_count = unit.status2.limbs_stand_max
    unit.status2.limbs_grasp_count = unit.status2.limbs_grasp_max
    unit.status2.limbs_fly_count = unit.status2.limbs_fly_max

    --print("Resetting status flags...")
    unit.flags2.has_breaks = false
    unit.flags2.gutted = false
    unit.flags2.vision_good = true
    unit.flags2.vision_damaged = false
    unit.flags2.vision_missing = false
    unit.flags2.breathing_good = true
    unit.flags2.breathing_problem = false

    unit.flags2.calculated_nerves = false
    unit.flags2.calculated_bodyparts = false
    unit.flags2.calculated_insulation = false
    unit.flags3.body_temp_in_range = false
    unit.flags3.compute_health = true
    unit.flags3.gelded = false

    --print("Resetting counters...")
    unit.counters.winded = 0
    unit.counters.stunned = 0
    unit.counters.unconscious = 0
    unit.counters.webbed = 0
    unit.counters.pain = 0
    unit.counters.nausea = 0
    unit.counters.dizziness = 0
    unit.counters.suffocation = 0
    unit.counters.guts_trail1.x = -30000
    unit.counters.guts_trail1.y = -30000
    unit.counters.guts_trail1.z = -30000
    unit.counters.guts_trail2.x = -30000
    unit.counters.guts_trail2.y = -30000
    unit.counters.guts_trail2.z = -30000

    unit.counters2.paralysis = 0
    unit.counters2.numbness = 0
    unit.counters2.fever = 0
    unit.counters2.exhaustion = 0
    unit.counters2.hunger_timer = 0
    unit.counters2.thirst_timer = 0
    unit.counters2.sleepiness_timer = 0
    unit.counters2.vomit_timeout = 0

    unit.animal.vanish_countdown = 0

    unit.body.infection_level = 0

    --print("Resetting body part status...")
    local comp = unit.body.components
    for i = 0, #comp.nonsolid_remaining - 1 do
        comp.nonsolid_remaining[i] = 100    -- percent remaining of fluid layers (Urist Da Vinci)
    end

    for i = 0, #comp.layer_wound_area - 1 do
        comp.layer_status[i].whole = 0        -- severed, leaking layers (Urist Da Vinci)
        comp.layer_wound_area[i] = 0        -- wound contact areas (Urist Da Vinci)
        comp.layer_cut_fraction[i] = 0        -- 100*surface percentage of cuts/fractures on the body part layer (Urist Da Vinci)
        comp.layer_dent_fraction[i] = 0        -- 100*surface percentage of dents on the body part layer (Urist Da Vinci)
        comp.layer_effect_fraction[i] = 0        -- 100*surface percentage of "effects" on the body part layer (Urist Da Vinci)
    end

    for _, status in ipairs(unit.body.components.body_part_status) do
        status.on_fire = false
        status.missing = false
        status.organ_loss = false
        status.organ_damage = false
        status.muscle_loss = false
        status.muscle_damage = false
        status.bone_loss = false
        status.bone_damage = false
        status.skin_damage = false
        status.motor_nerve_severed = false
        status.sensory_nerve_severed = false
        status.spilled_guts = false
        status.severed_or_jammed = false
    end

    unit.status2.body_part_temperature:resize(0) -- attempting to rewrite temperature was causing body parts to melt for some reason; forcing repopulation in this manner appears to be safer

    for i = 0,#unit.enemy.body_part_useable-1,1 do
        unit.enemy.body_part_useable[i] = 1 -- not sure what this does, but values appear to change following injuries
    end
    for i = 0,#unit.enemy.invorder_bp_start-1,1 do
        unit.enemy.invorder_bp_start[i] = 0 -- same as above
    end
    for i = 0,#unit.enemy.motor_nervenet-1,1 do
        unit.enemy.motor_nervenet[i] = 3 -- as above
    end
    for i = 0,#unit.enemy.sensory_nervenet-1,1 do
        unit.enemy.sensory_nervenet[i] = 3 -- as above
    end

    local histFig = df.historical_figure.find(unit.hist_figure_id)
    if histFig and histFig.info and histFig.info.wounds then
        --print("Clearing historical wounds...")
        histFig.info.wounds = nil
    end

    local health = unit.health
    if health then
        for i = 0, #health.flags-1,1 do
            health.flags[i] = false
        end
        for _,bpFlags in ipairs(health.body_part_flags) do
            for i = 0, #bpFlags-1,1 do
                bpFlags[i] = false
            end
        end
        health.immobilize_cntdn = 0
        health.dressing_cntdn = 0
        health.suture_cntdn = 0
        health.crutch_cntdn = 0
        health.try_for_cast_cntdn = 0
    end

    local job = unit.job.current_job
    if job and job.job_type == df.job_type.Rest then
        --print("Wake from rest...")
        job.completion_timer = 0
        job.pos:assign(unit.pos)
    end

    local job_link = df.global.world.jobs.list.next
    while job_link do
        local doctor_job = job_link.item
        if doctor_job then
            local patientRef = dfhack.job.getGeneralRef(doctor_job, df.general_ref_type['UNIT_PATIENT']) --as:df.general_ref_unit_patientst
            if patientRef and patientRef.unit_id == unit.id then
                patientRef.unit_id = -1 -- causes active healthcare job to be cancelled, generating a job cancellation announcement indicating the lack of a patient
                break
            end
        end
        job_link = job_link.next
    end
end

if dfhack_flags.module then
    return
end

if args.all then
    for _,unit in ipairs(df.global.world.units.active) do
        heal(unit,args.r,args.keep_corpse)
    end
elseif args.all_citizens then
    -- can't use dfhack.units.getCitizens since we want dead ones too
    for _,unit in ipairs(df.global.world.units.active) do
        if dfhack.units.isCitizen(unit, true) or dfhack.units.isResident(unit) then
            heal(unit,args.r,args.keep_corpse)
        end
    end
elseif args.all_civ then
    for _,unit in ipairs(df.global.world.units.active) do
        if dfhack.units.isOwnCiv(unit) then
            heal(unit,args.r,args.keep_corpse)
        end
    end
else
    local unit
    if args.unit then
        unit = df.unit.find(tonumber(args.unit))
        if not unit then
            qerror('Invalid unit ID: ' .. args.unit)
        end
    else
        local item = dfhack.gui.getSelectedItem(true)
        if item and df.item_corpsest:is_instance(item) then
            unit = df.unit.find(item.unit_id)
            if not unit then
                qerror('This corpse can no longer be resurrected.') -- unit has been offloaded
            end
            unit.pos:assign(xyz2pos(dfhack.items.getPosition(item))) -- to make the unit resurrect at the location of the corpse, rather than the location of death
        else
            unit = dfhack.gui.getSelectedUnit()
        end
    end
    if not unit then
        qerror('Please select a unit or corpse, or specify its ID via the -unit argument.')
    end
    heal(unit,args.r,args.keep_corpse)
end