2 files changed, 34 insertions, 19 deletions

diff --git a/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniItemTracker.kt b/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniItemTracker.kt
index 8d6f98d69..3cadd767b 100644
--- a/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniItemTracker.kt
+++ b/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniItemTracker.kt
@@ -25,8 +25,9 @@ class SkyHanniItemTracker<Data : ItemTrackerData>(
     name: String,
     createNewSession: () -> Data,
     getStorage: (ProfileSpecificStorage) -> Data,
+    vararg extraStorage: Pair<DisplayMode, (ProfileSpecificStorage) -> Data>,
     drawDisplay: (Data) -> List<Searchable>,
-) : SkyHanniTracker<Data>(name, createNewSession, getStorage, drawDisplay) {
+) : SkyHanniTracker<Data>(name, createNewSession, getStorage, *extraStorage, drawDisplay = drawDisplay) {
 
     companion object {
         val SKYBLOCK_COIN = NEUInternalName.SKYBLOCK_COIN
@@ -37,9 +38,8 @@ class SkyHanniItemTracker<Data : ItemTrackerData>(
             it.addItem(internalName, amount, command)
         }
         getSharedTracker()?.let { sharedData ->
-            sharedData.get(DisplayMode.TOTAL).items[internalName]?.let { data ->
-                sharedData.get(DisplayMode.SESSION).items[internalName]!!.hidden = data.hidden
-            }
+            val isHidden = sharedData.get(DisplayMode.TOTAL).items[internalName]?.hidden
+            if (isHidden != null) sharedData.modify { it.items[internalName]?.hidden = isHidden }
         }
 
         if (command) {
diff --git a/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniTracker.kt b/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniTracker.kt
index 5b07c9906..7810c1ee6 100644
--- a/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniTracker.kt
+++ b/src/main/java/at/hannibal2/skyhanni/utils/tracker/SkyHanniTracker.kt
@@ -6,6 +6,7 @@ import at.hannibal2.skyhanni.config.storage.ProfileSpecificStorage
 import at.hannibal2.skyhanni.data.ProfileStorageData
 import at.hannibal2.skyhanni.data.TrackerManager
 import at.hannibal2.skyhanni.features.misc.items.EstimatedItemValue
+import at.hannibal2.skyhanni.test.command.ErrorManager
 import at.hannibal2.skyhanni.utils.ChatUtils
 import at.hannibal2.skyhanni.utils.CollectionUtils
 import at.hannibal2.skyhanni.utils.NEUInternalName
@@ -24,9 +25,11 @@ open class SkyHanniTracker<Data : TrackerData>(
     val name: String,
     private val createNewSession: () -> Data,
     private val getStorage: (ProfileSpecificStorage) -> Data,
+    vararg extraStorage: Pair<DisplayMode, (ProfileSpecificStorage) -> Data>,
     private val drawDisplay: (Data) -> List<Searchable>,
 ) {
 
+    private val extraDisplayModes = extraStorage.toMap()
     private var inventoryOpen = false
     private var displayMode: DisplayMode? = null
     private val currentSessions = mutableMapOf<ProfileSpecificStorage, Data>()
@@ -61,12 +64,8 @@ open class SkyHanniTracker<Data : TrackerData>(
     }
 
     fun modify(mode: DisplayMode, modifyFunction: (Data) -> Unit) {
-        val storage = ProfileStorageData.profileSpecific ?: return
-        val data: Data = when (mode) {
-            DisplayMode.TOTAL -> storage.getTotal()
-            DisplayMode.SESSION -> storage.getCurrentSession()
-        }
-        modifyFunction(data)
+        val sharedTracker = getSharedTracker() ?: return
+        sharedTracker.modify(mode, modifyFunction)
         update()
     }
 
@@ -124,6 +123,8 @@ open class SkyHanniTracker<Data : TrackerData>(
         },
     )
 
+    private val availableTrackers = arrayOf(DisplayMode.TOTAL, DisplayMode.SESSION) + extraDisplayModes.keys
+
     private fun buildDisplayModeView() = Renderable.horizontalContainer(
         CollectionUtils.buildSelector<DisplayMode>(
             "§7Display Mode: ",
@@ -134,11 +135,17 @@ open class SkyHanniTracker<Data : TrackerData>(
                 storedTrackers[name] = it
                 update()
             },
+            universe = availableTrackers,
         ),
     )
 
-    protected fun getSharedTracker() = ProfileStorageData.profileSpecific?.let {
-        SharedTracker(it.getTotal(), it.getCurrentSession())
+    protected fun getSharedTracker() = ProfileStorageData.profileSpecific?.let { ps ->
+        SharedTracker(
+            mapOf(
+                DisplayMode.TOTAL to ps.getTotal(),
+                DisplayMode.SESSION to ps.getCurrentSession(),
+            ) + extraDisplayModes.mapValues { it.value(ps) },
+        )
     }
 
     private fun ProfileSpecificStorage.getCurrentSession() = currentSessions.getOrPut(this) { createNewSession() }
@@ -165,22 +172,30 @@ open class SkyHanniTracker<Data : TrackerData>(
         }
     }
 
-    class SharedTracker<Data : TrackerData>(private val total: Data, private val currentSession: Data) {
+    inner class SharedTracker<Data : TrackerData>(
+        private val entries: Map<DisplayMode, Data>,
+    ) {
 
-        fun modify(modifyFunction: (Data) -> Unit) {
-            modifyFunction(total)
-            modifyFunction(currentSession)
+        fun modify(mode: DisplayMode, modifyFunction: (Data) -> Unit) {
+            get(mode).let(modifyFunction)
         }
 
-        fun get(displayMode: DisplayMode) = when (displayMode) {
-            DisplayMode.TOTAL -> total
-            DisplayMode.SESSION -> currentSession
+        fun modify(modifyFunction: (Data) -> Unit) {
+            entries.values.forEach(modifyFunction)
         }
+
+        fun get(displayMode: DisplayMode) = entries[displayMode] ?: ErrorManager.skyHanniError(
+            "Unregistered display mode accessed on tracker",
+            "tracker" to name,
+            "displayMode" to displayMode,
+            "availableModes" to entries.keys,
+        )
     }
 
     enum class DisplayMode(val displayName: String) {
         TOTAL("Total"),
         SESSION("This Session"),
+        MAYOR("This Mayor"),
         ;
     }