targeting.lua

-- targeting.lua
-- Cluster-seeking AI override for high-volume (SMG/MG) and AOE (rocket/grenade/
-- flamethrower) weapons.  Each tick, pending ShootRangedWeapon actions for
-- qualifying weapons are redirected toward the densest enemy cluster rather than
-- DF's default nearest-threat targeting.
--
-- All mutable state lives in _G.__targeting so enable/disable/status/debug
-- commands issued as separate script invocations share the same runtime state.

local repeatUtil = require("repeat-util")
local eventful   = require("plugins.eventful")
local GCfg = _G.__GunConfig
if not GCfg or not GCfg._loaded then
    dfhack.printerr("[targeting] FATAL: __GunConfig not set — bw_autoload must run first")
    return
end

local CALLBACK_ID        = "targeting"
local CALLBACK_ID_RESCAN = "targeting:rescan"
local RESCAN_INTERVAL    = 100

-- Shared state
local S = _G.__targeting or {}
_G.__targeting = S

if S.enabled          == nil then S.enabled          = false end
if S.DEBUG            == nil then S.DEBUG            = false end
if S.cluster_carriers == nil then S.cluster_carriers = {} end
if S.redirect_count   == nil then S.redirect_count   = 0 end

---------------------------------------------------------------------------
-- CONFIG (loaded from config_gun.lua)
---------------------------------------------------------------------------

-- High-ROF sustained-fire weapons: redirect each shot to the density peak.
local SPRAY_WEAPONS = GCfg.SPRAY_WEAPONS
-- Single-shot AOE weapons: redirect to maximize enemies in blast radius.
local AOE_WEAPONS   = GCfg.AOE_WEAPONS
-- AOE blast radius per ammo type (Chebyshev tiles).
local AOE_AMMO_RADIUS = GCfg.AOE_AMMO_RADIUS
-- Penetrator ammo that does NOT explode; seek CREATURE_CLASS:ARMORED instead.
local ANTI_ARMOR_AMMO = GCfg.ANTI_ARMOR_AMMO
-- Anti-materiel rifle weapons: ARMORED priority targeting.
local AMR_WEAPONS   = GCfg.AMR_WEAPONS

-- Max enemy search range from the shooter (Chebyshev tiles).
local SEARCH_RANGE = 30
-- Spray: cluster radius and minimum group size to trigger redirect.
local SPRAY_CLUSTER_RADIUS = 2
local SPRAY_MIN_SIZE       = 3
-- AOE: minimum group size to trigger redirect.
local AOE_MIN_SIZE         = 2
-- Fallback AOE radius when ammo type is not in AOE_AMMO_RADIUS.
local DEFAULT_AOE_RADIUS   = 4

---------------------------------------------------------------------------
-- HELPERS
---------------------------------------------------------------------------

local function get_ammo_subtype_id(item)
    local ok, id = pcall(function()
        return dfhack.items.getSubtypeDef(
            df.item_type.AMMO, item:getSubtype()).id
    end)
    return ok and id or nil
end

local function chebyshev(a, b)
    return math.max(math.abs(a.x - b.x), math.abs(a.y - b.y))
end

-- Returns true if unit_b is an ally of unit_a (should not be targeted).
local function are_allies(unit_a, unit_b)
    if unit_a.id == unit_b.id then return false end
    if unit_a.civ_id ~= -1 and unit_a.civ_id == unit_b.civ_id then return true end
    if unit_a.military.squad_id ~= -1 and unit_a.military.squad_id == unit_b.military.squad_id then return true end
    if dfhack.units.isCitizen(unit_a) and dfhack.units.isCitizen(unit_b) then return true end
    return false
end

-- Returns the weapon subtype ID string of the ranged weapon a unit is carrying.
local function get_unit_weapon_id(unit)
    local result = nil
    pcall(function()
        for _, inv_item in ipairs(unit.inventory) do
            local item = inv_item.item
            if item:getType() == df.item_type.WEAPON then
                local ok, wid = pcall(function()
                    return dfhack.items.getSubtypeDef(
                        df.item_type.WEAPON, item:getSubtype()).id
                end)
                if ok and wid then
                    result = wid
                    return
                end
            end
        end
    end)
    return result
end

-- Returns the position of the enemy whose neighbourhood contains the most
-- other enemies, or nil if no cluster meets the minimum threshold.
local function find_cluster_target(shooter, cluster_radius, min_size)
    local enemies = {}
    for _, u in ipairs(df.global.world.units.active) do
        if u.id == shooter.id then goto next_unit end  -- never target own position
        if dfhack.units.isActive(u) and not dfhack.units.isDead(u)
                and not are_allies(shooter, u)
                and math.abs(u.pos.z - shooter.pos.z) <= 1
                and chebyshev(u.pos, shooter.pos) <= SEARCH_RANGE then
            table.insert(enemies, u)
        end
        ::next_unit::
    end

    if #enemies < min_size then return nil end

    local best_pos, best_count = nil, min_size - 1
    for _, e in ipairs(enemies) do
        local count = 0
        for _, e2 in ipairs(enemies) do
            if chebyshev(e.pos, e2.pos) <= cluster_radius then
                count = count + 1
            end
        end
        if count > best_count then
            best_count = count
            best_pos   = e.pos
        end
    end
    return best_pos
end

-- Returns true if the unit's race (or caste) has the given creature class string.
local function has_creature_class(unit, class_name)
    local ok, result = pcall(function()
        local race = df.global.world.raws.creatures.all[unit.race]
        if not race then return false end
        for _, cc in ipairs(race.creature_class) do
            if cc == class_name then return true end
        end
        local caste = race.caste[unit.caste]
        if caste then
            for _, cc in ipairs(caste.creature_class) do
                if cc == class_name then return true end
            end
        end
        return false
    end)
    return ok and result or false
end

-- Returns the position of the nearest enemy with CREATURE_CLASS:ARMORED,
-- or nil if none exist within SEARCH_RANGE.
local function find_armored_target(shooter)
    local best_pos, best_dist = nil, SEARCH_RANGE + 1
    for _, u in ipairs(df.global.world.units.active) do
        if u.id == shooter.id then goto next_unit end
        if dfhack.units.isActive(u) and not dfhack.units.isDead(u)
                and not are_allies(shooter, u)
                and math.abs(u.pos.z - shooter.pos.z) <= 1
                and has_creature_class(u, "ARMORED") then
            local dist = chebyshev(u.pos, shooter.pos)
            if dist < best_dist then
                best_dist = dist
                best_pos  = u.pos
            end
        end
        ::next_unit::
    end
    return best_pos
end

-- Returns true if the unit has any anti-armor ammo in inventory.
local function unit_carries_anti_armor_ammo(unit)
    for _, inv_item in ipairs(unit.inventory) do
        local item = inv_item.item
        if item:getType() == df.item_type.AMMO then
            local sub_id = get_ammo_subtype_id(item)
            if sub_id and ANTI_ARMOR_AMMO[sub_id] then return true end
        end
    end
    return false
end

---------------------------------------------------------------------------
-- CARRIER MAINTENANCE
---------------------------------------------------------------------------

local function unit_is_cluster_carrier(unit)
    local wid = get_unit_weapon_id(unit)
    if wid ~= nil then
        if SPRAY_WEAPONS[wid] or AOE_WEAPONS[wid] then return true end
        if AMR_WEAPONS[wid] then return true end
    end
    -- Also track units with anti-armor ammo (e.g. rocket launcher with ROCKET2_UNPRIMED).
    return unit_carries_anti_armor_ammo(unit)
end

local function refresh_carriers()
    S.cluster_carriers = {}
    local adv = dfhack.world.getAdventurer()
    for _, unit in ipairs(df.global.world.units.active) do
        if adv and adv.id == unit.id then goto continue end
        if not dfhack.units.isActive(unit) then goto continue end
        if dfhack.units.isDead(unit) then goto continue end
        if unit_is_cluster_carrier(unit) then
            S.cluster_carriers[unit.id] = true
        end
        ::continue::
    end
end

local function on_inventory_change(unit_id, item_id, old_equip, new_equip)
    if not S.enabled then return end
    local unit = df.unit.find(unit_id)
    if not unit then
        S.cluster_carriers[unit_id] = nil
        return
    end
    local adv = dfhack.world.getAdventurer()
    if adv and adv.id == unit.id then return end

    if unit_is_cluster_carrier(unit) then
        S.cluster_carriers[unit_id] = true
    else
        S.cluster_carriers[unit_id] = nil
    end
end

local function rescan_carriers()
    if not S.enabled then return end
    refresh_carriers()
end

---------------------------------------------------------------------------
-- TICK HANDLER
---------------------------------------------------------------------------

local function check_carriers()
    if not S.enabled then return end

    for unit_id in pairs(S.cluster_carriers) do
        local unit = df.unit.find(unit_id)
        if not unit then
            S.cluster_carriers[unit_id] = nil
            goto continue
        end

        for i = #unit.actions - 1, 0, -1 do
            local action = unit.actions[i]
            if action.type ~= df.unit_action_type.ShootRangedWeapon then
                goto next_action
            end

            pcall(function()
                local d = action.data.shootrangedweapon
                local weapon_id = get_unit_weapon_id(unit)
                if not weapon_id then return end

                local cluster_pos = nil

                if AOE_WEAPONS[weapon_id] then
                    local ammo = df.item.find(d.ammo_itid)
                    local sub_id = ammo and get_ammo_subtype_id(ammo) or nil

                    if sub_id and ANTI_ARMOR_AMMO[sub_id] then
                        -- Non-exploding penetrator: seek ARMORED priority, not clusters.
                        cluster_pos = find_armored_target(unit)
                    else
                        local radius = DEFAULT_AOE_RADIUS
                        if sub_id and AOE_AMMO_RADIUS[sub_id] then
                            radius = AOE_AMMO_RADIUS[sub_id]
                        end
                        cluster_pos = find_cluster_target(unit, radius, AOE_MIN_SIZE)
                    end
                elseif SPRAY_WEAPONS[weapon_id] then
                    cluster_pos = find_cluster_target(unit, SPRAY_CLUSTER_RADIUS, SPRAY_MIN_SIZE)
                else
                    -- AMR weapon or other anti-armor ammo carrier: ARMORED priority only.
                    local ammo = df.item.find(d.ammo_itid)
                    local sub_id = ammo and get_ammo_subtype_id(ammo) or nil
                    if sub_id and ANTI_ARMOR_AMMO[sub_id] then
                        cluster_pos = find_armored_target(unit)
                    end
                end

                if cluster_pos then
                    d.target_lx = cluster_pos.x
                    d.target_ly = cluster_pos.y
                    d.target_lz = cluster_pos.z
                    S.redirect_count = S.redirect_count + 1
                    if S.DEBUG then
                        dfhack.print(("[targeting] unit %d (%s) redirected to cluster at (%d,%d,%d)\n"):format(
                            unit_id, weapon_id, cluster_pos.x, cluster_pos.y, cluster_pos.z))
                    end
                end
            end)

            ::next_action::
        end

        ::continue::
    end
end

---------------------------------------------------------------------------
-- ENABLE / DISABLE / STATUS / DEBUG
---------------------------------------------------------------------------

local function enable()
    if S.enabled then
        print("[targeting] Already enabled")
        return
    end
    S.enabled = true
    S.redirect_count = 0
    S.cluster_carriers = {}
    refresh_carriers()
    eventful.enableEvent(eventful.eventType.INVENTORY_CHANGE, 5)
    eventful.onInventoryChange[CALLBACK_ID] = on_inventory_change
    repeatUtil.scheduleEvery(CALLBACK_ID, 1, "ticks", check_carriers)
    repeatUtil.scheduleEvery(CALLBACK_ID_RESCAN, RESCAN_INTERVAL, "ticks", rescan_carriers)
    print("[targeting] Enabled")
end

local function disable()
    if not S.enabled then
        print("[targeting] Already disabled")
        return
    end
    S.enabled = false
    S.cluster_carriers = {}
    eventful.onInventoryChange[CALLBACK_ID] = nil
    repeatUtil.cancel(CALLBACK_ID)
    repeatUtil.cancel(CALLBACK_ID_RESCAN)
    print("[targeting] Disabled")
end

local function status()
    local count = 0
    for _ in pairs(S.cluster_carriers) do count = count + 1 end
    print(("[targeting] Status: %s | Carriers: %d | Redirects since enable: %d"):format(
        S.enabled and "ENABLED" or "DISABLED",
        count,
        S.redirect_count))
end

local function toggle_debug()
    S.DEBUG = not S.DEBUG
    print(("[targeting] Debug %s"):format(S.DEBUG and "ON" or "OFF"))
end

---------------------------------------------------------------------------
-- ARGUMENT PARSING
---------------------------------------------------------------------------

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

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