1 files changed, 56 insertions, 5 deletions

diff --git a/src/main/kotlin/moe/nea/firmament/features/inventory/SlotLocking.kt b/src/main/kotlin/moe/nea/firmament/features/inventory/SlotLocking.kt
index 76a719a..943f7dd 100644
--- a/src/main/kotlin/moe/nea/firmament/features/inventory/SlotLocking.kt
+++ b/src/main/kotlin/moe/nea/firmament/features/inventory/SlotLocking.kt
@@ -3,23 +3,30 @@
  *
  * SPDX-License-Identifier: GPL-3.0-or-later
  */
+@file:UseSerializers(DashlessUUIDSerializer::class)
 
 package moe.nea.firmament.features.inventory
 
+import java.util.*
 import org.lwjgl.glfw.GLFW
 import kotlinx.serialization.Serializable
+import kotlinx.serialization.UseSerializers
 import kotlinx.serialization.serializer
 import net.minecraft.entity.player.PlayerInventory
+import net.minecraft.screen.slot.SlotActionType
 import moe.nea.firmament.events.HandledScreenKeyPressedEvent
 import moe.nea.firmament.events.IsSlotProtectedEvent
 import moe.nea.firmament.events.SlotRenderEvents
 import moe.nea.firmament.features.FirmamentFeature
 import moe.nea.firmament.gui.config.ManagedConfig
+import moe.nea.firmament.keybindings.SavedKeyBinding
 import moe.nea.firmament.mixins.accessor.AccessorHandledScreen
 import moe.nea.firmament.util.CommonSoundEffects
 import moe.nea.firmament.util.MC
 import moe.nea.firmament.util.SBData
 import moe.nea.firmament.util.data.ProfileSpecificDataHolder
+import moe.nea.firmament.util.json.DashlessUUIDSerializer
+import moe.nea.firmament.util.skyblockUUID
 
 object SlotLocking : FirmamentFeature {
     override val identifier: String
@@ -29,10 +36,15 @@ object SlotLocking : FirmamentFeature {
     data class Data(
         val lockedSlots: MutableSet<Int> = mutableSetOf(),
         val lockedSlotsRift: MutableSet<Int> = mutableSetOf(),
+
+        val lockedUUIDs: MutableSet<UUID> = mutableSetOf(),
     )
 
     object TConfig : ManagedConfig(identifier) {
-        val lock by keyBinding("lock") { GLFW.GLFW_KEY_L }
+        val lockSlot by keyBinding("lock") { GLFW.GLFW_KEY_L }
+        val lockUUID by keyBindingWithOutDefaultModifiers("lock-uuid") {
+            SavedKeyBinding(GLFW.GLFW_KEY_L, shift = true)
+        }
     }
 
     override val config: TConfig
@@ -40,6 +52,8 @@ object SlotLocking : FirmamentFeature {
 
     object DConfig : ProfileSpecificDataHolder<Data>(serializer(), "locked-slots", ::Data)
 
+    val lockedUUIDs get() = DConfig.data?.lockedUUIDs
+
     val lockedSlots
         get() = when (SBData.skyblockLocation) {
             "rift" -> DConfig.data?.lockedSlotsRift
@@ -49,7 +63,7 @@ object SlotLocking : FirmamentFeature {
 
     override fun onLoad() {
         HandledScreenKeyPressedEvent.subscribe {
-            if (!it.matches(TConfig.lock)) return@subscribe
+            if (!it.matches(TConfig.lockSlot)) return@subscribe
             val inventory = MC.handledScreen ?: return@subscribe
             inventory as AccessorHandledScreen
 
@@ -65,19 +79,56 @@ object SlotLocking : FirmamentFeature {
                 CommonSoundEffects.playSuccess()
             }
         }
+        HandledScreenKeyPressedEvent.subscribe {
+            if (!it.matches(TConfig.lockUUID)) return@subscribe
+            val inventory = MC.handledScreen ?: return@subscribe
+            inventory as AccessorHandledScreen
+
+            val slot = inventory.focusedSlot_Firmament ?: return@subscribe
+            val stack = slot.stack ?: return@subscribe
+            val uuid = stack.skyblockUUID ?: return@subscribe
+            val lockedUUIDs = lockedUUIDs ?: return@subscribe
+            if (uuid in lockedUUIDs) {
+                lockedUUIDs.remove(uuid)
+            } else {
+                lockedUUIDs.add(uuid)
+            }
+            DConfig.markDirty()
+            CommonSoundEffects.playSuccess()
+        }
+        IsSlotProtectedEvent.subscribe {
+            if (it.slot != null && it.slot.inventory is PlayerInventory && it.slot.index in (lockedSlots ?: setOf())) {
+                it.protect()
+            }
+        }
         IsSlotProtectedEvent.subscribe {
-            if (it.slot.inventory is PlayerInventory && it.slot.index in (lockedSlots ?: setOf())) {
+            if (it.actionType == SlotActionType.SWAP
+                || it.actionType == SlotActionType.PICKUP
+                || it.actionType == SlotActionType.QUICK_MOVE
+                || it.actionType == SlotActionType.QUICK_CRAFT
+                || it.actionType == SlotActionType.CLONE
+                || it.actionType == SlotActionType.PICKUP_ALL
+            ) return@subscribe
+            val stack = it.itemStack ?: return@subscribe
+            val uuid = stack.skyblockUUID ?: return@subscribe
+            if (uuid in (lockedUUIDs ?: return@subscribe)) {
                 it.protect()
             }
         }
         SlotRenderEvents.Before.subscribe {
-            if (it.slot.inventory is PlayerInventory && it.slot.index in (lockedSlots ?: setOf())) {
+            val isSlotLocked = it.slot.inventory is PlayerInventory && it.slot.index in (lockedSlots ?: setOf())
+            val isUUIDLocked = (it.slot.stack?.skyblockUUID ?: return@subscribe) in (lockedUUIDs ?: setOf())
+            if (isSlotLocked || isUUIDLocked) {
                 it.context.fill(
                     it.slot.x,
                     it.slot.y,
                     it.slot.x + 16,
                     it.slot.y + 16,
-                    0xFFFF0000.toInt()
+                    when {
+                        isSlotLocked -> 0xFFFF0000.toInt()
+                        isUUIDLocked -> 0xFF00FF00.toInt()
+                        else -> error("Slot is locked, but not by slot or uuid")
+                    }
                 )
             }
         }