exterminate.lua

--@module = true

local argparse = require('argparse')

local function spawnLiquid(position, liquid_level, liquid_type, update_liquids)
  local map_block = dfhack.maps.getTileBlock(position)
  local tile = dfhack.maps.getTileFlags(position)

  tile.flow_size = liquid_level
  tile.liquid_type = liquid_type
  tile.flow_forbid = false

  map_block.flags.update_liquid = update_liquids
  map_block.flags.update_liquid_twice = update_liquids
end

local function checkUnit(unit)
    return (unit.body.blood_count ~= 0 or unit.body.blood_max == 0) and
        not unit.flags1.inactive and
        not unit.flags1.caged and
        not unit.flags1.chained
end

local function isUnitFriendly(unit)
    if dfhack.units.isDanger(unit) then
        return false
    end
    local adv = dfhack.world.getAdventurer()
    if adv then
        if adv == unit or
            unit.relationship_ids.GroupLeader == adv.id or
            unit.relationship_ids.PetOwner == adv.id
        then
            return true
        end
    end

    return dfhack.units.isOwnCiv(unit) or
        dfhack.units.isOwnGroup(unit) or
        dfhack.units.isVisiting(unit) or
        dfhack.units.isTame(unit) or
        dfhack.units.isDomesticated(unit)
end

killMethod = {
    INSTANT = 0,
    BUTCHER = 1,
    MAGMA = 2,
    DROWN = 3,
    VAPORIZE = 4,
    DISINTEGRATE = 5,
    KNOCKOUT = 6,
    TRAUMATIZE = 7,
}

-- removes the unit from existence, leaving no corpse if the unit hasn't died
-- by the time the vanish countdown expires
local function vaporizeUnit(unit, target_value)
    target_value = target_value or 1
    unit.animal.vanish_countdown = target_value
end

-- Kills a unit by removing blood and also setting a vanish countdown as a failsafe.
local function destroyUnit(unit)
    unit.body.blood_count = 0
    vaporizeUnit(unit, 2)
end

--  Marks a unit for slaughter at the butcher's shop.
local function butcherUnit(unit)
    unit.flags2.slaughter = true
end

--  Knocks a unit out for 30k ticks or the target value
local function knockoutUnit(unit, target_value)
    target_value = target_value or 30000
    unit.counters.unconscious = target_value
end

--  Traumatizes the unit, forcing them to stare off into space. Cuts down on pathfinding
local function traumatizeUnit(unit)
    unit.mood = df.mood_type.Traumatized
end


local function drownUnit(unit, liquid_type)
    previousPositions = previousPositions or {}
    previousPositions[unit.id] = copyall(unit.pos)

    local function createLiquid()
        spawnLiquid(unit.pos, 7, liquid_type)
        if not same_xyz(previousPositions[unit.id], unit.pos) then
            spawnLiquid(previousPositions[unit.id], 0, nil, false)
            previousPositions[unit.id] = copyall(unit.pos)
        end
        if unit.flags2.killed then
            spawnLiquid(previousPositions[unit.id], 0, nil, false)
        else
            dfhack.timeout(1, 'ticks', createLiquid)
        end
    end
    createLiquid()
end

local function destroyInventory(unit)
    for index = #unit.inventory-1, 0, -1 do
        local item = unit.inventory[index].item
        dfhack.items.remove(item)
    end
end

function killUnit(unit, method)
    if method == killMethod.BUTCHER then
        butcherUnit(unit)
    elseif method == killMethod.MAGMA then
        drownUnit(unit, df.tile_liquid.Magma)
    elseif method == killMethod.DROWN then
        drownUnit(unit, df.tile_liquid.Water)
    elseif method == killMethod.VAPORIZE then
        vaporizeUnit(unit)
    elseif method == killMethod.DISINTEGRATE then
        vaporizeUnit(unit)
        destroyInventory(unit)
    elseif method == killMethod.KNOCKOUT then
        knockoutUnit(unit)
    elseif method == killMethod.TRAUMATIZE then
        traumatizeUnit(unit)
    else
        destroyUnit(unit)
    end
end

local function getRaceCastes(race_id)
    local unit_castes = {}
    for _, caste in pairs(df.creature_raw.find(race_id).caste) do
        unit_castes[caste.caste_id] = {}
    end
    return unit_castes
end

local function getMapRaces(only_visible, include_friendly)
    local map_races = {}
    for _, unit in pairs(df.global.world.units.active) do
        if only_visible and not dfhack.units.isVisible(unit) then
            goto skipunit
        end
        if not include_friendly and isUnitFriendly(unit) then
            goto skipunit
        end
        if dfhack.units.isActive(unit) and checkUnit(unit) then
            local unit_race_name = dfhack.units.isUndead(unit) and "UNDEAD" or df.creature_raw.find(unit.race).creature_id

            local race = ensure_key(map_races, unit_race_name)
            race.id = unit.race
            race.name = unit_race_name
            race.count = (race.count or 0) + 1
        end
        :: skipunit ::
    end

    return map_races
end

if dfhack_flags.module then
    return
end

local options, args = {
    help = false,
    method = killMethod.INSTANT,
    only_visible = false,
    include_friendly = false,
    limit = -1,
}, {...}

local positionals = argparse.processArgsGetopt(args, {
    {'h', 'help', handler = function() options.help = true end},
    {'m', 'method', handler = function(arg) options.method = killMethod[arg:upper()] end, hasArg = true},
    {'o', 'only-visible', handler = function() options.only_visible = true end},
    {'f', 'include-friendly', handler = function() options.include_friendly = true end},
    {'l', 'limit', handler = function(arg) options.limit = argparse.positiveInt(arg, 'limit') end, hasArg = true},
})

if not dfhack.isMapLoaded() then
    qerror('This script requires a fortress map to be loaded')
end

if positionals[1] == "help" or options.help then
    print(dfhack.script_help())
    return
end

if positionals[1] == "this" then
    local selected_unit = dfhack.gui.getSelectedUnit()
    if not selected_unit then
        qerror("Select a unit and run the script again.")
    end
    killUnit(selected_unit, options.method)
    print('Unit exterminated.')
    return
end

local map_races = getMapRaces(options.only_visible, options.include_friendly)

if not positionals[1] then
    local sorted_races = {}
    for race, value in pairs(map_races) do
        table.insert(sorted_races, { name = race, count = value.count })
    end
    table.sort(sorted_races, function(a, b)
        return a.count > b.count
    end)
    for _, race in ipairs(sorted_races) do
        print(([[%4s %s]]):format(race.count, race.name))
    end
    return
end

local count, target = 0, 'creature(s)'
local race_name = table.concat(positionals, ' ')
if race_name:lower() == 'undead' then
    target = 'undead'
    if not map_races.UNDEAD then
        qerror("No undead found on the map.")
    end
    for _, unit in pairs(df.global.world.units.active) do
        if dfhack.units.isUndead(unit) and checkUnit(unit) then
            killUnit(unit, options.method)
            count = count + 1
        end
    end
elseif positionals[1]:split(':')[1] == "all" then
    local selected_caste = positionals[1]:split(':')[2]

    for _, unit in ipairs(df.global.world.units.active) do
        if options.limit > 0 and count >= options.limit then
            break
        end
        if not checkUnit(unit) then
            goto skipunit
        end
        if options.only_visible and not dfhack.units.isVisible(unit) then
            goto skipunit
        end
        if not options.include_friendly and isUnitFriendly(unit) then
            goto skipunit
        end
        if selected_caste and selected_caste ~= df.creature_raw.find(unit.race).caste[unit.caste].caste_id then
            goto skipunit
        end

        killUnit(unit, options.method)
        count = count + 1
        :: skipunit ::
    end
else
    local selected_race, selected_caste = race_name, nil

    if string.find(selected_race, ':') then
        local tokens = selected_race:split(':')
        selected_race, selected_caste = tokens[1], tokens[2]
    end

    if not map_races[selected_race] then
        local selected_race_upper = selected_race:upper()
        local selected_race_under = selected_race_upper:gsub(' ', '_')
        if map_races[selected_race_upper] then
            selected_race = selected_race_upper
        elseif map_races[selected_race_under] then
            selected_race = selected_race_under
        else
            qerror("No creatures of this race on the map (" .. selected_race .. ").")
        end
    end

    local race_castes = getRaceCastes(map_races[selected_race].id)

    if selected_caste and not race_castes[selected_caste] then
        local selected_caste_upper = selected_caste:upper()
        if race_castes[selected_caste_upper] then
            selected_caste = selected_caste_upper
        else
            qerror("Invalid caste: " .. selected_caste)
        end
    end

    target = selected_race

    for _, unit in pairs(df.global.world.units.active) do
        if options.limit > 0 and count >= options.limit then
            break
        end
        if not checkUnit(unit) then
            goto skipunit
        end
        if options.only_visible and not dfhack.units.isVisible(unit) then
            goto skipunit
        end
        if not options.include_friendly and isUnitFriendly(unit) then
            goto skipunit
        end
        if selected_caste and selected_caste ~= df.creature_raw.find(unit.race).caste[unit.caste].caste_id then
            goto skipunit
        end

        if selected_race == df.creature_raw.find(unit.race).creature_id then
            killUnit(unit, options.method)
            count = count + 1
        end

        :: skipunit ::
    end
end

print(([[Exterminated %d %s.]]):format(count, target))