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")
LuaUserData.RegisterType("Barotrauma.Store")
LuaUserData.RegisterType("Barotrauma.GUIComponent")
LuaUserData.RegisterType("Barotrauma.PurchasedItem")
LuaUserData.RegisterType("Barotrauma.ItemPrefab")
LuaUserData.RegisterType("Barotrauma.Location+StoreInfo")
LuaUserData.MakeMethodAccessible(Descriptors["Barotrauma.CargoManager"], "GetConfirmedSoldEntities")

-- Simple configuration
local CONFIG = {
    QUICKSTACK_KEYS = Keys.F,
    FABRICATOR_KEY = Keys.V,
    MAX_BUY = Keys.B,
    NESTED_CONTAINERS = true,
    DEBUG_MODE = true,
}

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

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

---@param table table
---@param depth number?
local function DumpTable(table, depth)
    if depth == nil then
        depth = 0
    end
    if (depth > 200) then
        print("Error: Depth > 200 in dumpTable()")
        return
    end
    for k, v in pairs(table) do
        if (type(v) == "table") then
            print(string.rep("  ", depth) .. k .. ":")
            DumpTable(v, depth + 1)
        else
            print(string.rep("  ", depth) .. k .. ": ", v)
        end
    end
end

-- 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(string.format("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(string.format("Building item tree for inventory at slot index: %d", slotIndex))
        -- debugPrint(string.format("Slot %d has %d items", slotIndex, #slot.items))
        if #slot.items == 0 then
            -- debugPrint(string.format("Slot %d is empty, adding to itemTree as 'empty'", slotIndex))
            itemTree['empty'] = itemTree['empty'] or {}
            itemTree['empty'][#itemTree['empty'] + 1] = {
                inventory = inventory,
                slotIndex = slotIndex - 1,
                maxFits = 60
            }
            -- debugPrint(string.format("Added empty slot to itemTree at index: %d", slotIndex))
        else
            ---@type Barotrauma.Item
            local item = slot.items[1]
            local identifier = item.Prefab.Identifier.Value
            -- debugPrint(string.format("Found item: %s with identifier: %s", item.Name, 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(string.format("Added item to itemTree under identifier: %s", 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(string.format("Searching inside %s for nested containers", item.Name))
                buildItemTree(item.OwnInventory, itemTree)
            end
        end
    end

    -- debugPrint("Completed building item tree")
    debugPrint("Completed building item tree")
    return itemTree
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

-- We got to do this shit because enqueueInventory calls enqueueItem
-- And enqueueItem calls enqueueInventory
-- So unless we define them both before using them
-- We will get an error saying either is undefined
local enqueueItem
local enqueueSlot
local enqueueInventory
local _

---@param item Barotrauma.Item
---@param queue Barotrauma.Item[]
---@param predicate? fun(item: Barotrauma.Item): boolean
---@return Barotrauma.Item[], string
enqueueItem = function(item, queue, predicate)
    queue = queue or {}
    predicate = predicate or function() return true end
    -- debugPrint(string.format("Enqueuing item: %s", item.Prefab.Identifier.Value))
    -- local err
    if item.OwnInventory then
        -- As far as I know every item has only one inventory
        -- Only machines have multiple
        -- So inventrorY should be fine here
        -- debugPrint("Item has its own inventory, enqueuing inventory...")
        queue, _ = enqueueInventory(item.OwnInventory, queue)
        -- if err then
        --     debugPrint(string.format("Error enqueuing inventory: %s", err))
        -- end
    end
    if predicate(item) then queue[#queue + 1] = item end
    -- debugPrint(string.format("Item enqueued. Current queue size: %d", #queue))
    return queue, nil
end

---@param slot Barotrauma.ItemInventory.Slot
---@param queue Barotrauma.Item[]
---@param predicate? fun(item: Barotrauma.Item): boolean
---@return Barotrauma.Item[], string
enqueueSlot = function(slot, queue, predicate)
    queue = queue or {}
    predicate = predicate or function() return true end
    -- debugPrint(string.format("Enqueuing slot with %d items.", #slot.items))
    -- 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
        -- Only the final leaf nodes decide upon the predicate
        queue, _ = enqueueItem(item, queue, predicate)
        -- if err then
        --     debugPrint(string.format("Error enqueuing item: %s", err))
        -- end
    end
    -- debugPrint(string.format("Finished enqueuing slot. Current queue size: %d", #queue))
    return queue
end

---@param inventory Barotrauma.ItemInventory
---@param queue Barotrauma.Item[]
---@param predicate? fun(item: Barotrauma.Item): boolean
---@return Barotrauma.Item[]
enqueueInventory = function(inventory, queue, predicate)
    queue = queue or {}
    predicate = predicate or function() return true end
    -- debugPrint(string.format("Enqueuing inventory with %d slots.", #inventory.slots))
    -- local err
    for _, slot in ipairs(inventory.slots) do
        -- Only the final leaf nodes decide upon the predicate
        queue, _ = enqueueSlot(slot, queue, predicate)
        -- if err then
        --     debugPrint(string.format("Error enqueuing slot: %s", err))
        -- end
    end
    -- debugPrint(string.format("Finished enqueuing inventory. Current queue size: %d", #queue))
    return queue
end

-- This is a bit fucking sucky.....
-- But I really don't know better
-- Maybe it will be fine...
---@return Barotrauma.Item[]
local function getOpenContainers()
    debugPrint("Attempting to find open container...")
    -- local containers = {}
    -- for item in Item.ItemList do
    --     ---@cast item Barotrauma.Item
    --     local isok = true
    --     isok = isok and item ~= nil
    --     isok = isok and item.OwnInventory ~= nil
    --     isok = isok and item.OwnInventory.visualSlots ~= nil
    --     isok = isok and #item.OwnInventory.visualSlots > 0
    --     -- I don't know what rootContainer is
    --     -- It seems to be the parent of the current item...?
    --     -- Maybe the world object...
    --     -- Either way - static objects that we may open have it
    --     -- And our own inventory does not
    --     -- So it's a good selector for now
    --     isok = isok and item.rootContainer ~= nil
    --     if isok then
    --         containers[#containers + 1] = item
    --     end
    -- end

    local controlledCharacter = Character.Controlled
    if not controlledCharacter then return {} end
    local selectedItem = controlledCharacter.SelectedItem
    if not selectedItem then return {} end
    return { selectedItem }
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(string.format("Preparing to stack items into the bag..."))
    local bagSlot = inventory.slots[8]
    if bagSlot then
        -- debugPrint(string.format("Bag slot found at index 8 with %d items.", #bagSlot.items))
        if #bagSlot.items > 0 then
            local item = bagSlot.items[1]
            -- debugPrint(string.format("Found item in bag slot: %s", item.Name))
            if item and item.OwnInventory then
                -- debugPrint(string.format("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
-- 6 and 7 are hands
-- 9..18 are main slots
local inventorySlotsToStack = { 6, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18 }
local BAG_SLOT = 8
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(string.format("Error building item tree: %s", err))
        return
    end
    itemTree = sortItemtreeBySlots(itemTree)
    local toMove = {}

    -- for i, slot in ipairs(inventory.slots) do
    --     if #slot.items > 0 then
    --         local item = slot.items[1]
    --         local identifier = item.Prefab.Identifier.Value
    --         print(string.format("Item at slot %d is %s", i, identifier))
    --     end
    -- end

    for _, slotid in ipairs(inventorySlotsToStack) do
        debugPrint(string.format("Processing inventory slot: %d", slotid))
        local slot = inventory.slots[slotid]
        if #slot.items > 0 then
            local item = slot.items[1]
            local tags = item.Prefab.Tags
            local shouldSuss = true
            for tag in tags do
                if tag.value:find("tool") or tag.value:find("weapon") then
                    debugPrint(string.format("Item '%s' is a tool or weapon, skipping", item.Name))
                    shouldSuss = false
                    break
                end
            end

            if shouldSuss then
                local before = #toMove
                toMove = enqueueSlot(slot, toMove)
                local after = #toMove
                debugPrint(string.format("Enqueued %d items from the inventory slot %d", after - before, slotid))
            end
        end
    end

    local openContainers = getOpenContainers()
    for _, container in ipairs(openContainers) do
        local inventories = container.OwnInventories
        debugPrint(string.format("Found %d inventories in the open container", #inventories))
        for i, containerInventory in ipairs(inventories) do
            debugPrint(string.format("Enqueuing inventory %d with %d slots", i, #containerInventory.slots))
            local before = #toMove
            toMove = enqueueInventory(containerInventory, toMove)
            local after = #toMove
            debugPrint(string.format("Enqueued %d items from the open container", after - before))
        end
    end

    local errors = tryMoveItems(toMove, itemTree)
    for _, error in ipairs(errors) do
        print(string.format("Error stacking item: %s", error))
    end
end

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

    local character = instance
    if not character then return end

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

----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------

---@return {item: Barotrauma.Item, fabricator: Barotrauma.FabricatorComponent}, string?
local function getOpenFabricator()
    -- Get the controlled character
    local controlledCharacter = Character.Controlled
    if not controlledCharacter then return nil, "No controlled character found" end

    -- Check if the character has a selected item
    local selectedItem = controlledCharacter.SelectedItem
    if not selectedItem then return nil, "No selected item found" end

    -- Check if the selected item has a Fabricator component
    local fabricator = Game.GetFabricatorComponent(selectedItem)
    if not fabricator then return nil, "No fabricator component found" end

    return {
        item = selectedItem,
        fabricator = fabricator
    }
end

--- Recipes can have multiple inputs, for example ammo
--- Can be made either out of copper iron or steel, 1 of either
---@class RecipeInfo
---@field targetItem {identifier: string, name: string, amount: number}
---@field requiredItems {amount: number, minCondition: number, maxCondition: number, prefabs: string[]}[]

---@param fabricator Barotrauma.FabricatorComponent
---@return RecipeInfo, string?
local function getSelectedRecipeRequirements(fabricator)
    -- local openFabricator, err = getOpenFabricator()
    -- if err then return nil, err end
    -- local fabricator = openFabricator.fabricator

    local selectedRecipe = fabricator.SelectedItem
    if not selectedRecipe then return nil, "No selected recipe found" end

    local requiredItems = {}
    for _, requiredItem in pairs(selectedRecipe.RequiredItems) do
        local itemInfo = {
            amount = tonumber(requiredItem.Amount),
            minCondition = tonumber(requiredItem.MinCondition),
            maxCondition = tonumber(requiredItem.MaxCondition),
            prefabs = {}
        }

        for prefab in requiredItem.ItemPrefabs do
            itemInfo.prefabs[#itemInfo.prefabs + 1] = prefab.Identifier
        end

        requiredItems[#requiredItems + 1] = itemInfo
    end

    return {
        targetItem = {
            identifier = selectedRecipe.TargetItem.Identifier,
            name = selectedRecipe.TargetItem.Name,
            amount = selectedRecipe.Amount
        },
        requiredItems = requiredItems
    }
end

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

    -- TODO: Maybe get items from entire sub...?
    -- There's no point getting recipes if we don't have all of this bullshit
    ---@type Barotrauma.Character
    local character = instance
    if not character then
        debugPrint("Character instance is nil.")
        return
    end
    ---@type Barotrauma.CharacterInventory
    local inventory = character.Inventory
    if not inventory then
        debugPrint("Character inventory is nil.")
        return
    end
    ---@type Barotrauma.ItemInventory.Slot
    local bagSlot = inventory.slots[BAG_SLOT]
    if not bagSlot then
        debugPrint("Bag slot not found.")
        return
    end
    if #bagSlot.items == 0 then
        debugPrint("Bag slot is empty.")
        return
    end
    ---@type Barotrauma.Item
    local bagItem = bagSlot.items[1]
    if not bagItem then
        debugPrint("Bag item not found.")
        return
    end

    local fabricator, err = getOpenFabricator()
    if err then
        print(string.format("Error getting open fabricator: %s", err))
        return
    end

    local recipe
    recipe, err = getSelectedRecipeRequirements(fabricator.fabricator)
    if err then
        print(string.format("Error getting selected recipe requirements: %s", err))
    end
    DumpTable(recipe)

    -- TODO: Maybe make it so every press cycles the input
    -- For recipes that have multiple prefabs
    -- But then again what if it has 3 items with 4 prefabs each..
    -- Is that 4 iterations or 3*4 iterations?
    local toFind = recipe.requiredItems
    ---@type Barotrauma.Item[]
    local toGet = {}
    ---@type fun(item: Barotrauma.Item): boolean
    local filter = function(item)
        local found = false
        if #toFind == 0 then return false end
        -- toFind are all items we need to find
        for i, itemInfo in ipairs(toFind) do
            -- prefabs are all items that satisfy the requirements
            for _, prefab in ipairs(itemInfo.prefabs) do
                if item.Prefab.Identifier == prefab then
                    toGet[#toGet + 1] = item
                    itemInfo.amount = itemInfo.amount - 1
                    found = true
                    break
                end
            end
            if itemInfo.amount <= 0 then
                toFind[i] = nil
            end
            if found then break end
        end
        return found
    end

    local items = enqueueInventory(bagItem.OwnInventory, {}, filter)
    -- TODO: This might explode... Oh well?
    local inputInventory = fabricator.item.OwnInventories[1]
    local slot = -1
    local previous = nil
    for _, item in ipairs(items) do
        if previous ~= item.Prefab.Identifier then slot = slot + 1 end
        inputInventory.TryPutItem(item, slot, false, true, nil)
        previous = item.Prefab.Identifier
    end
    DumpTable(items)
end, Hook.HookMethodType.After)

----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------


---@return Barotrauma.Location.StoreInfo[], string?
local function getCurrentStore()
    if not Game or not Game.GameSession or not Game.GameSession.Campaign then
        return nil, "No game session found"
    end

    local map = Game.GameSession.Campaign.Map
    if not map or not map.CurrentLocation or not map.CurrentLocation.Stores then
        return nil, "No map found"
    end

    local location = map.CurrentLocation

    -- Otherwise, determine which store is active by checking the cargo manager
    local cargoManager = Game.GameSession.Campaign.CargoManager
    if not cargoManager then
        return nil, "No cargo manager found"
    end

    -- Find which store has items in the cart
    local stores = {}
    for _, store in pairs(location.Stores) do
        if #cargoManager:GetBuyCrateItems(store) > 0 then
            stores[#stores + 1] = store
        end
    end

    return stores, nil
end

-- Example: Add a key binding to buy all items in the current store
-- when the 'B' key is pressed
Hook.Patch("Barotrauma.Character", "ControlLocalPlayer", function(instance, ptable)
    if PlayerInput.KeyHit(CONFIG.MAX_BUY) then
        local cargoManager = Game.GameSession.Campaign.CargoManager
        if not cargoManager then
            print("No cargo manager available")
            return
        end

        local stores, err = getCurrentStore()
        if err then
            print(string.format("Error getting current store: %s", err))
            return
        end

        for _, store in ipairs(stores) do
            local toAdd = {}
            -- Get items available at the store
            local items = cargoManager:GetBuyCrateItems(store)
            for item in items do
                -- We have already added this many of item
                toAdd[item.ItemPrefab.Identifier.Value] = -item.Quantity
            end
            for item in store.Stock do
                -- So if we add the total amount available
                -- We get the amount we have to add to buy entire stock
                if toAdd[item.ItemPrefab.Identifier.Value] then
                    toAdd[item.ItemPrefab.Identifier.Value] = toAdd[item.ItemPrefab.Identifier.Value] + item.Quantity
                end
            end
            for item, amount in pairs(toAdd) do
                if amount > 0 then
                    print(string.format("Adding %d of %s to the buy crate", amount, item))
                    cargoManager.ModifyItemQuantityInBuyCrate(store.Identifier, item, amount)
                end
            end
        end
    end
end, Hook.HookMethodType.After)