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

---@param item Barotrauma.Item
---@param itemTree table<string, ItemLocation[]>
---@return string
local function tryMoveItem(item, itemTree)
    local location = itemTree[item.Prefab.Identifier.Value]
    if not location then return nil, "No locations for item, not stacking" end

    local moved = false
    -- First try to move to existing stacks
    for _, itemLocation in ipairs(location) do
        if itemLocation.maxFits > 0 then
            moved = moved or itemLocation.inventory.TryPutItem(item, itemLocation.slotIndex, false, true, nil)
            itemLocation.maxFits = itemLocation.inventory.HowManyCanBePut(item.Prefab, itemLocation.slotIndex)
        end
    end

    -- If we can not find an existing stack
    -- Then move to any of the empty slots
    if not moved then
        for _, itemLocation in ipairs(itemTree['empty']) do
            moved = moved or itemLocation.inventory.TryPutItem(item, itemLocation.slotIndex, false, true, nil)
            itemLocation.maxFits = itemLocation.inventory.HowManyCanBePut(item.Prefab, itemLocation.slotIndex)
        end
    end

    -- If we still can not move the item give up
    if not moved then return "Failed to find valid location for item" end
    return nil
end

---@param items Barotrauma.Item[]
---@param itemTree table<string, ItemLocation[]>
---@return string[]
local function tryMoveItems(items, itemTree)
    local errs = {}
    for _, item in ipairs(items) do
        local err = tryMoveItem(item, itemTree)
        -- oops, this one failed, continue...
        if err then
            errs[#errs + 1] = string.format("Failed to move item: %s", item.Prefab.Identifier.Value)
        end
    end
    return errs
end

---@param item Barotrauma.Item
---@param queue Barotrauma.Item[]
---@return Barotrauma.Item[], string
local function enqueueItem(item, queue)
    queue = queue or {}
    queue[#queue + 1] = item
    return queue
end

---@param slot Barotrauma.ItemInventory.Slot
---@param queue Barotrauma.Item[]
---@return Barotrauma.Item[], string
local function enqueueSlot(slot, queue)
    queue = queue or {}
    -- We don't want to shadow queue
    local err
    -- If the slot is empty there's nothing to iterate
    -- And we will naturally return queue as is
    for _, item in ipairs(slot.items) do
        queue, err = enqueueItem(item, queue)
        if err then
            print("Error enqueuing item: " .. err)
        end
    end
    return queue
end

---@param inventory Barotrauma.ItemInventory
---@param queue Barotrauma.Item[]
---@return Barotrauma.Item[], string[]
local function enqueueInventory(inventory, queue)
    queue = queue or {}
    -- We don't want to shadow queue
    local err
    for _, slot in ipairs(inventory.slots) do
        queue, err = enqueueSlot(slot, queue)
        if err then
            print("Error enqueuing slot: " .. err)
        end
    end
    return queue
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

local function stackPlayerInventoryItems(inventory, itemTree)
    debugPrint("Starting to stack player 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

---@param inventory Barotrauma.ItemInventory
---@return table<string, ItemLocation[]>, string
local function tryBuildItemTree(inventory)
    local 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
                return itemTree, "Bag does not have its own inventory."
            end
        else
            return itemTree, "Bag slot is empty."
        end
    else
        return itemTree, "No bag slot found at index 8."
    end
    return itemTree, nil
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, err = tryBuildItemTree(inventory)
    if err then
        debugPrint("Error building item tree: " .. err)
        return
    end
    itemTree = sortItemtreeBySlots(itemTree)

    local toMove = enqueueInventory(inventory)
    for _, item in ipairs(toMove) do
        print("Item: " .. item.Prefab.Identifier.Value)
    end
    local errors = tryMoveItems(toMove, itemTree)
    for _, error in ipairs(errors) do
        print("Error stacking item: " .. error)
    end

    -- stackPlayerInventoryItems(inventory, 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)