Update to make catagories

common/media/lua/shared/OpinionatedFirearms_AmmoMakerSpentCasingsDisplayCategoryPatch.lua

@@ -0,0 +1,127 @@
+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