OpinionatedFirearms/common/media/lua/shared/OpinionatedFirearms_AmmoMakerSpentCasingsDisplayCategoryPatch.lua

local OFAmmoMakerDisplayCategoryPatch = {
    patched = false,
    tickHookAdded = false,
}

local SPENT_CASINGS_CATEGORY = "OFSpentCasings"
local AMMO_PART_CATEGORY = "OFAmmoPart"

local function addAmmoMakerType(typeSet, fullType)
    if type(fullType) ~= "string" or fullType == "" then
        return
    end

    if not string.find(fullType, "^ammomaker%.") then
        return
    end

    typeSet[fullType] = true
end

local function setDisplayCategory(scriptManager, fullType, category)
    local scriptItem = scriptManager:FindItem(fullType)
    if not scriptItem then
        return false
    end

    if type(scriptItem.DoParam) == "function" then
        scriptItem:DoParam("DisplayCategory = " .. category)
        return true
    end

    if type(scriptItem.setDisplayCategory) == "function" then
        scriptItem:setDisplayCategory(category)
        return true
    end

    return false
end

local function collectAmmoMakerTypes()
    if type(ammoMakerAmmoParts) ~= "table" then
        return nil, nil
    end

    local spentCasingTypes = {}
    local ammoPartTypes = {}

    for _, partData in pairs(ammoMakerAmmoParts) do
        if type(partData) == "table" then
            local partClass = partData.partClass
            if partClass == "Casing" or partClass == "Hull" then
                addAmmoMakerType(spentCasingTypes, partData.partOld)
                addAmmoMakerType(spentCasingTypes, partData.partFired)
                addAmmoMakerType(ammoPartTypes, partData.boxType)
                addAmmoMakerType(ammoPartTypes, partData.bagType)
            end
        end
    end

    return spentCasingTypes, ammoPartTypes
end

local function applyCategory(scriptManager, itemTypes, category)
    local patchedAny = false

    for fullType in pairs(itemTypes) do
        if setDisplayCategory(scriptManager, fullType, category) then
            patchedAny = true
        end
    end

    return patchedAny
end

local function applyPatch()
    if OFAmmoMakerDisplayCategoryPatch.patched then
        return true
    end

    local scriptManager = ScriptManager and ScriptManager.instance
    if not scriptManager or type(scriptManager.FindItem) ~= "function" then
        return false
    end

    local spentCasingTypes, ammoPartTypes = collectAmmoMakerTypes()
    if not spentCasingTypes or not ammoPartTypes then
        return false
    end

    local patchedSpentCasings = applyCategory(scriptManager, spentCasingTypes, SPENT_CASINGS_CATEGORY)
    local patchedAmmoParts = applyCategory(scriptManager, ammoPartTypes, AMMO_PART_CATEGORY)
    if not patchedSpentCasings and not patchedAmmoParts then
        return false
    end

    OFAmmoMakerDisplayCategoryPatch.patched = true
    return true
end

local function tryPatchOnTick()
    if applyPatch() then
        OFAmmoMakerDisplayCategoryPatch.tickHookAdded = false
        Events.OnTick.Remove(tryPatchOnTick)
    end
end

local function ensureTickHook()
    if OFAmmoMakerDisplayCategoryPatch.tickHookAdded then
        return
    end

    OFAmmoMakerDisplayCategoryPatch.tickHookAdded = true
    Events.OnTick.Add(tryPatchOnTick)
end

if not applyPatch() then
    ensureTickHook()
end

if Events.OnMainMenuEnter and type(Events.OnMainMenuEnter.Add) == "function" then
    Events.OnMainMenuEnter.Add(function()
        OFAmmoMakerDisplayCategoryPatch.patched = false
        if not applyPatch() then
            ensureTickHook()
        end
    end)
end