armaments.lua

-- armaments.lua
-- Auto-equips specific creatures/castes with hand-targeted weapons/ammo.
-- Current config: HUMAN MALE and HUMAN FEMALE get bow in right hand and arrows in left.

local eventful = require("plugins.eventful")

local CALLBACK_ID = "armaments"

local enabled = false
local processed_units = {} -- [unit_id] = true

local MODE_WEAPON = df.inv_item_role_type.Weapon

-- Easy-to-extend loadout sets keyed by RACE:CASTE.
-- Each set supports multiple paired options; one is chosen randomly per unit.
local LOADOUTS = {
    ["VEHICLE:TANK_LIGHT"] = {
        options = {
            {
                right = { id = "WEAPON:ITEM_WEAPON_BOW_LONG",  material = "INORGANIC:BASE_METAL" },
                left  = { id = "AMMO:ITEM_AMMO_ARROW",         material = "INORGANIC:METAL_FERROUS" },
            },
            {
                right = { id = "WEAPON:ITEM_WEAPON_CROSSBOW",  material = "INORGANIC:BASE_METAL" },
                left  = { id = "AMMO:ITEM_AMMO_BOLT_LIGHT",    material = "INORGANIC:METAL_FERROUS" },
            },
        },
    },
    ["VEHICLE:TANK_HEAVY"] = {
        options = {
            {
                right = { id = "WEAPON:ITEM_WEAPON_BOW_LONG",  material = "INORGANIC:BASE_METAL" },
                left  = { id = "AMMO:ITEM_AMMO_ARROW",         material = "INORGANIC:METAL_FERROUS" },
            },
            {
                right = { id = "WEAPON:ITEM_WEAPON_CROSSBOW",  material = "INORGANIC:BASE_METAL" },
                left  = { id = "AMMO:ITEM_AMMO_BOLT_LIGHT",    material = "INORGANIC:METAL_FERROUS" },
            },
        },
    },
}

local function get_race_token(unit)
    local race = df.creature_raw.find(unit.race)
    return race and race.creature_id or nil
end

local function get_caste_token(unit)
    local race = df.creature_raw.find(unit.race)
    if not race then return nil end
    local caste = race.caste[unit.caste]
    return caste and caste.caste_id or nil
end

local function has_inventory(unit)
    return unit.inventory and #unit.inventory > 0
end

local function find_hand_parts(unit)
    local right_hand = -1
    local left_hand = -1
    local grasp_parts = {}

    if not unit.body or not unit.body.body_plan then
        return right_hand, left_hand
    end

    for bp_idx, bp in ipairs(unit.body.body_plan.body_parts) do
        if bp.flags.GRASP then
            table.insert(grasp_parts, bp_idx)
            if bp.flags.RIGHT and right_hand < 0 then
                right_hand = bp_idx
            elseif bp.flags.LEFT and left_hand < 0 then
                left_hand = bp_idx
            end
        end
    end

    if right_hand < 0 and #grasp_parts >= 1 then right_hand = grasp_parts[1] end
    if left_hand < 0 then
        if #grasp_parts >= 2 then
            left_hand = grasp_parts[2]
        elseif #grasp_parts >= 1 then
            left_hand = grasp_parts[1]
        end
    end

    return right_hand, left_hand
end

local function create_item_for_unit(unit, item_id, material)
    local item_type = dfhack.items.findType(item_id)
    if item_type < 0 then
        qerror(("[armaments] Unknown item type in '%s'"):format(item_id))
    end
    local item_subtype = dfhack.items.findSubtype(item_id)

    local mat = dfhack.matinfo.find(material)
    if not mat then
        qerror(("[armaments] Unknown material '%s'"):format(material))
    end

    local items = dfhack.items.createItem(unit, item_type, item_subtype, mat.type, mat.index)
    return items and items[1] or nil
end

local function equip_loadout(unit, loadout)
    local right_hand, left_hand = find_hand_parts(unit)
    if right_hand < 0 or left_hand < 0 then
        dfhack.printerr(("[armaments] Unit %d has no valid hand body parts"):format(unit.id))
        return false
    end

    local right_item = create_item_for_unit(unit, loadout.right.id, loadout.right.material)
    if not right_item then
        dfhack.printerr(("[armaments] Failed to create right-hand item for unit %d"):format(unit.id))
        return false
    end

    local left_item = create_item_for_unit(unit, loadout.left.id, loadout.left.material)
    if not left_item then
        dfhack.printerr(("[armaments] Failed to create left-hand item for unit %d"):format(unit.id))
        return false
    end

    local ok_right = dfhack.items.moveToInventory(right_item, unit, MODE_WEAPON, right_hand)
    local ok_left = dfhack.items.moveToInventory(left_item, unit, MODE_WEAPON, left_hand)

    if not ok_right then
        dfhack.printerr(("[armaments] Failed to place right-hand item for unit %d"):format(unit.id))
    end
    if not ok_left then
        dfhack.printerr(("[armaments] Failed to place left-hand item for unit %d"):format(unit.id))
    end

    return ok_right and ok_left
end

local function get_loadout_for_unit(unit)
    local race_token = get_race_token(unit)
    local caste_token = get_caste_token(unit)
    if not race_token or not caste_token then return nil, nil end
    local key = race_token .. ":" .. caste_token
    return LOADOUTS[key], key
end

local function pick_option(loadout)
    if not loadout or type(loadout.options) ~= "table" or #loadout.options == 0 then
        return nil
    end
    return loadout.options[math.random(#loadout.options)]
end

local function should_arm(unit)
    if not unit then return false end
    if not dfhack.units.isAlive(unit) then return false end
    if has_inventory(unit) then return false end
    if processed_units[unit.id] then return false end

    local loadout = get_loadout_for_unit(unit)
    return loadout ~= nil
end

local function try_arm_unit(unit)
    if not should_arm(unit) then return end
    local loadout, key = get_loadout_for_unit(unit)
    if not loadout then return end
    local option = pick_option(loadout)
    if not option then return end
    if equip_loadout(unit, option) then
        processed_units[unit.id] = true
        print(("[armaments] Armed %s unit %d"):format(key, unit.id))
    end
end

local function scan_units()
    local units = df.global.world.units.active
    for i = 0, #units - 1 do
        try_arm_unit(units[i])
    end
end

local function on_unit_new_active(unit_id)
    if not enabled then return end
    local unit = df.unit.find(unit_id)
    if not unit then return end
    try_arm_unit(unit)
end

local function enable()
    if enabled then
        print("[armaments] Already enabled")
        return
    end

    processed_units = {}
    enabled = true

    scan_units()
    eventful.enableEvent(eventful.eventType.UNIT_NEW_ACTIVE, 1)
    eventful.onUnitNewActive[CALLBACK_ID] = on_unit_new_active

    print("[armaments] Enabled - startup scan complete")
end

local function disable()
    if not enabled then
        print("[armaments] Already disabled")
        return
    end

    eventful.onUnitNewActive[CALLBACK_ID] = nil
    enabled = false
    processed_units = {}
    print("[armaments] Disabled")
end

local function status()
    local n = 0
    for _ in pairs(processed_units) do n = n + 1 end
    print(("[armaments] Status: %s"):format(enabled and "ENABLED" or "DISABLED"))
    print(("[armaments] Processed units: %d"):format(n))
    print("[armaments] Configured loadout sets:")
    for key, loadout in pairs(LOADOUTS) do
        local option_count = (loadout and loadout.options and #loadout.options) or 0
        print(("[armaments]   %s (%d options)"):format(key, option_count))
    end
end

local args = { ... }
local command = args[1] or "enable"

if command == "enable" then
    enable()
elseif command == "disable" then
    disable()
elseif command == "status" then
    status()
else
    print("[armaments] Usage: armaments [enable|disable|status]")
end