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)
    PRIORITY_CONTAINERS = { "toolbelt", "backpack", "pouch" }, -- Priority order for containers
    EXCLUDE_TOOLS = true,                                      -- Exclude tools from container list
    TOOL_IDENTIFIERS = { "welding", "gun", "weapon", "revolver", "smg", "rifle", "shotgun",
        "diving", "oxygen", "scanner", "card", "id", "fuel", "rod", "battery",
        "fabricator", "deconstructor" },                 -- Common tool identifiers to exclude
    ALWAYS_INCLUDE = { "toolbelt", "backpack", "pouch" } -- Always include these items even if they match tool criteria
}

-- 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, iter)
    iter = iter or 0
    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
        -- This iteration is done only to skip the player inventory as a potential destination
        -- Since it will always be 0th iteration and there's no point stacking
        -- Items from the player inventory to itself
        debugPrint("Building item tree for inventory at iteration: " .. iter .. ", slot index: " .. slotIndex)
        debugPrint("Slot " .. slotIndex .. " has " .. #slot.items .. " items")
        if iter > 0 then
            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)
            end
        end


        if #slot.items > 0 then
            ---@type Barotrauma.Item
            local item = slot.items[1]
            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, iter + 1)
            end
        end
    end

    debugPrint("Completed building item tree for current inventory iteration: " .. iter)
    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

            ---@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 = buildItemTree(inventory, {})
    itemTree = sortItemtreeBySlots(itemTree)
    stackInventoryItems(inventory, itemTree)

    local openContainerInventory = getOpenContainer()
    if openContainerInventory then
        stackInventoryItems(openContainerInventory, itemTree)
    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)