if SERVER then return end

-- Register necessary types and make fields accessible
LuaUserData.RegisterType("Barotrauma.Items.Components.ItemContainer+SlotRestrictions")
LuaUserData.RegisterType(
    'System.Collections.Immutable.ImmutableArray`1[[Barotrauma.Items.Components.ItemContainer+SlotRestrictions, Barotrauma]]')
LuaUserData.MakeFieldAccessible(Descriptors['Barotrauma.Items.Components.ItemContainer'], 'slotRestrictions')
LuaUserData.MakeFieldAccessible(Descriptors['Barotrauma.ItemInventory'], 'slots')
LuaUserData.MakeFieldAccessible(Descriptors["Barotrauma.CharacterInventory"], "slots")

-- Simple configuration
local CONFIG = {
    TRIGGER_KEY = Keys.F,     -- Key to press for quick stacking
    NESTED_CONTAINERS = true, -- Whether to include nested containers
    DEBUG_MODE = true,        -- Print debug messages
    HAND_SLOTS = { 6, 7 },    -- Slot numbers for hands (typically slots 6 and 7)
}

-- MOD INFO
local MOD_NAME = "Quick Stack To Containers"
local MOD_VERSION = "1.1.0"

print(MOD_NAME .. " v" .. MOD_VERSION .. " loaded!")

-- Debugging helper function
local function debugPrint(message)
    if CONFIG.DEBUG_MODE then
        print("[" .. MOD_NAME .. "] " .. message)
    end
end

---@class ItemLocation
---@field inventory Barotrauma.ItemInventory
---@field slotIndex number
---@field maxFits number

-- The resulting item tree is a table where the key is an ID of an item
-- And the value is an object that represents where that item is located
-- In our inventory
-- Special case are empty slots where any item fits
---@param inventory Barotrauma.ItemInventory
---@param itemTree table<string, ItemLocation[]>
---@return table
local function buildItemTree(inventory, itemTree)
    itemTree = itemTree or {}
    if not inventory or not inventory.slots then
        debugPrint("Inventory is nil or has no slots, returning empty itemTree")
        return itemTree
    end

    -- One slot can have one item but multiple of it
    -- The number of an item in a slot is #slot.items
    for slotIndex, slot in ipairs(inventory.slots) do
        debugPrint("Building item tree for inventory at slot index: " .. slotIndex)
        debugPrint("Slot " .. slotIndex .. " has " .. #slot.items .. " items")
        if #slot.items == 0 then
            debugPrint("Slot " .. slotIndex .. " is empty, adding to itemTree as 'empty'")
            itemTree['empty'] = itemTree['empty'] or {}
            itemTree['empty'][#itemTree['empty'] + 1] = {
                inventory = inventory,
                slotIndex = slotIndex - 1,
                maxFits = 60
            }
            debugPrint("Added empty slot to itemTree at index: " .. slotIndex)
        else
            ---@type Barotrauma.Item
            local item = slot.items[1]
            local identifier = item.Prefab.Identifier.Value
            debugPrint("Found item: " .. item.Name .. " with identifier: " .. identifier)
            itemTree[identifier] = itemTree[identifier] or {}
            -- We DO want even slots with maxFits = 0
            -- Because that indicates that we DO HAVE the item
            -- At all
            -- And based on that we decide to move it
            itemTree[identifier][#itemTree[identifier] + 1] = {
                inventory = inventory,
                slotIndex = slotIndex - 1,
                maxFits = slot.HowManyCanBePut(item.Prefab)
            }
            debugPrint("Added item to itemTree under identifier: " .. identifier)

            local tags = item.Prefab.Tags
            local shouldSuss = false
            for tag in tags do
                if tag.value:find("container") then
                    shouldSuss = true
                    break
                end
            end

            if shouldSuss then
                debugPrint("Searching inside " .. item.Name .. " for nested containers")
                buildItemTree(item.OwnInventory, itemTree)
            end
        end
    end

    debugPrint("Completed building item tree")
    return itemTree
end

---@param slot Barotrauma.ItemInventory.Slot
---@param itemLocation ItemLocation
local function moveItemsTo(slot, itemLocation)
    local totalHere = #slot.items
    local moveHere = math.min(itemLocation.maxFits, totalHere)
    debugPrint("Attempting to move " ..
        moveHere .. " items to inventory at slot index: " .. itemLocation.slotIndex)
    if moveHere > 0 then
        for _, item in ipairs(slot.items) do
            itemLocation.inventory.TryPutItem(item, itemLocation.slotIndex, false, true, nil)
            itemLocation.maxFits = itemLocation.inventory.HowManyCanBePut(item.Prefab, itemLocation.slotIndex)
            if itemLocation.maxFits <= 0 then return end
        end
    end
end

local function stackInventoryItems(inventory, itemTree)
    debugPrint("Starting to stack inventory items...")
    for slotIndex, slot in ipairs(inventory.slots) do
        debugPrint("Checking slot index: " .. slotIndex)
        -- Cannot stack items if there are no items...
        if #slot.items > 0 then
            ---@type Barotrauma.Item
            local item = slot.items[1]
            local identifier = item.Prefab.Identifier.Value
            debugPrint("Item at slot " .. slotIndex .. " is " .. identifier)

            ---@type ItemLocation[]
            local locations = itemTree[identifier]
            -- If there are no locations for this item
            -- Then there's nowhere to move it
            if locations then
                for _, location in ipairs(locations) do
                    moveItemsTo(slot, location)
                    if #slot.items == 0 then break end
                end

                -- If we have processed all the locations and we still have items to move
                -- Then put them into the empty slots:
                if #slot.items > 0 then
                    for _, location in ipairs(itemTree['empty']) do
                        moveItemsTo(slot, location)
                        if #slot.items == 0 then break end
                    end
                end
            end
        else
            debugPrint("Slot index " .. slotIndex .. " is empty.")
        end
    end
    debugPrint("Completed stacking inventory items.")
end

-- This is a bit fucking sucky.....
-- But I really don't know better
-- Maybe it will be fine...
---@return Barotrauma.ItemInventory
local function getOpenContainer()
    debugPrint("Attempting to find open container...")
    for item in Item.ItemList do
        ---@cast item Barotrauma.Item
        if item and item.OwnInventory then
            if item.OwnInventory.visualSlots and #item.OwnInventory.visualSlots > 0 then
                return item.OwnInventory
            end
        end
    end

    debugPrint("No open container found")
    return nil
end

-- We would like to fill larger stacks first
---@param itemTree table<string, ItemLocation[]>
---@return table<string, ItemLocation[]>
local function sortItemtreeBySlots(itemTree)
    for _, item in pairs(itemTree) do
        table.sort(item, function(a, b)
            ---@cast a ItemLocation
            ---@cast b ItemLocation
            return a.maxFits < b.maxFits
        end)
    end

    return itemTree
end

-- Function to quickly stack items from inventory to containers
local function quickStackItems(character)
    if not character then
        debugPrint("No character found")
        return
    end

    debugPrint("Quick stack function called")

    local inventory = character.Inventory
    if not inventory or not inventory.slots then
        debugPrint("Character has no inventory")
        return
    end

    local itemTree = {}
    itemTree = buildItemTree(inventory, itemTree)
    itemTree = sortItemtreeBySlots(itemTree)

    debugPrint("Preparing to stack items into the bag...")
    local bagSlot = inventory.slots[8]
    if bagSlot then
        debugPrint("Bag slot found at index 8 with " .. #bagSlot.items .. " items.")
        if #bagSlot.items > 0 then
            local item = bagSlot.items[1]
            debugPrint("Found item in bag slot: " .. item.Name)
            if item and item.OwnInventory then
                debugPrint("Item has its own inventory, building item tree for it...")
                itemTree = buildItemTree(item.OwnInventory, itemTree)
            else
                debugPrint("Item does not have its own inventory.")
            end
        else
            debugPrint("Bag slot is empty.")
        end
    else
        debugPrint("No bag slot found at index 8.")
    end

    stackInventoryItems(item.OwnInventory, itemTree)
    local openContainerInventory = getOpenContainer()
    if openContainerInventory then
        stackInventoryItems(openContainerInventory, itemTree)
    end

    --local handItems = {}
    --for _, slotIndex in ipairs(CONFIG.HAND_SLOTS) do
    --    local slot = inventory.slots[slotIndex]
    --    debugPrint("Checking hand slot index: " .. slotIndex)
    --    if slot then
    --        debugPrint("Hand slot " .. slotIndex .. " found with " .. #slot.items .. " items.")
    --        if #slot.items > 0 then
    --            handItems[#handItems + 1] = slot.items[1]
    --            debugPrint("Added item " .. slot.items[1].Name .. " from hand slot " .. slotIndex .. " to handItems.")
    --        else
    --            debugPrint("Hand slot " .. slotIndex .. " is empty.")
    --        end
    --    else
    --        debugPrint("Hand slot " .. slotIndex .. " does not exist.")
    --    end
    --end
end

-- Hook into player control to listen for key press
Hook.Patch("Barotrauma.Character", "ControlLocalPlayer", function(instance, ptable)
    if not PlayerInput.KeyHit(CONFIG.TRIGGER_KEY) then return end

    local character = instance
    if not character then return end

    quickStackItems(character)
end, Hook.HookMethodType.After)