cc-utilities/storage/itemgroup.lua

local Sentinel = require("storage.sentinel")

local ItemGroup = Sentinel:tag({}, Sentinel.ITEMGROUP)
ItemGroup.__index = ItemGroup

--- Create a group of managed itemstacks
---@param stack table ItemStack to compare to
---@return table ItemGroup instance representing stack type
function ItemGroup:from(stack)
    local obj = { stack }
    setmetatable(obj, self)
    obj.__index = obj
    return obj
end

function ItemGroup.collectStacks(stacks)
  local groups = {}
  for _,stack in ipairs(stacks) do
    for _,group in ipairs(groups) do
      if group:addStack(stack) then
        goto continue
      end
    end
    table.insert(groups, ItemGroup:from(stack))
      ::continue::
  end
  return groups
end

function ItemGroup:getStackCount()
  return #self
end

function ItemGroup:getDisplayName()
  return self:_getIdentityStack():getDisplayName()
end

function ItemGroup:getSimpleName()
  return self:isEmpty() and "[EMPTY]" or self:_getIdentityStack():getSimpleName()
end

function ItemGroup:isEmpty()
  return self:_getIdentityStack():isEmpty()
end

function ItemGroup:getName()
  return self:_getIdentityStack():getName()
end

function ItemGroup:getEnchantments()
  return self:_getIdentityStack():getEnchantments()
end

function ItemGroup:getNBT()
  return self:_getIdentityStack():getNBT()
end

function ItemGroup:getDamage()
  return self:_getIdentityStack():getDamage()
end

function ItemGroup:getMaxDamage()
  return self:_getIdentityStack():getMaxDamage()
end

function ItemGroup:_getIdentityStack()
  return self[1]
end

function ItemGroup:_iterateStacks()
  return ipairs(self)
end

function ItemGroup:_addStack(stack)
  table.insert(self, stack)
end

function ItemGroup:_removeStack(stackOrIndex)
  if type(stackOrIndex) == "number" then
    return table.remove(self, stackOrIndex) ~= nil
  else
    for index,stack in self:_iterateStacks() do
      if stack == stackOrIndex then
        return self:_removeStack(index)
      end
    end
  end
end

function ItemGroup:canAddStack(stack)
  if self:_getIdentityStack():canTransfer(stack) then
    for _,chkStack in self:_iterateStacks() do
        if chkStack == stack then
            -- Already in the group
            return true
        end
    end
    return true
  end
  return false
end

function ItemGroup:addStack(stack)
  if self:canAddStack(stack) then
    self:_addStack(stack)
    return true
  end
  return false
end

function ItemGroup:transferTo(target, itemCount)
  local targetGroup = nil
  if Sentinel:is(target, Sentinel.CHEST) or Sentinel:is(target, Sentinel.STORAGE) then
    local identity = self:_getIdentityStack()
    local find = target:find(function(stack) return identity:canTransfer(stack) end)
    if #find == 0 then
      return itemCount == nil or itemCount == 0, 0
    end

    targetGroup = ItemGroup:from(find[1])
    for i=2,#find do
      targetGroup:_addStack(find[i])
    end
  elseif Sentinel:is(target, Sentinel.ITEMSTACK) then
    targetGroup = ItemGroup:from(target)
  elseif Sentinel:is(target, Sentinel.ITEMGROUP) then
    targetGroup = target
  end

  if targetGroup == nil then
    error("Unexpected transfer target for ItemGroup:transferTo")
  end

  local targetCap = 0
  for _,stack in targetGroup:_iterateStacks() do
    targetCap = targetCap + (stack:getMaxCount() - stack:getCount())
  end

  -- TODO: Not efficient
  local transferMax = math.min(itemCount or targetCap, targetCap, self:getItemCount())
  local transfer = 0
  for _,stack in targetGroup:_iterateStacks() do
    for _,from in self:_iterateStacks() do
      if transfer >= transferMax then
        goto complete
      end

      if stack:getCount() == stack:getMaxCount() then
        goto continue
      end
      local _, xfer = from:transferTo(stack, math.min(stack:getMaxCount() - stack:getCount(), transferMax - transfer))
      transfer = transfer + xfer
    end
      ::continue::
  end
    ::complete::

  return itemCount == nil or (itemCount == transfer), transfer
end

function ItemGroup:getItemCount()
  if self:_getIdentityStack():isEmpty() then
    return 0
  end
  local sum = 0
  for _,stack in self:_iterateStacks() do
    sum = sum + stack:getCount()
  end
  return sum
end

function ItemGroup:_copyTrackedStackList()
  local collect = {}
  for index,stack in self:_iterateStacks() do
    table.insert(collect, { stack = stack, index = index })
  end
  return collect
end

function ItemGroup:defragment()
  if self:_getIdentityStack():isEmpty() then
    return 0
  end

  local entries = self:_copyStackList()
  table.sort(entries, function(a, b)
    return a.stack:getCount() > b.stack:getCount()
  end)
  local endPtr = #entries
  local startPtr = endPtr
  for i=1,endPtr do
    local entry = entries[i].stack
    if entry:getMaxCount() ~= entry:getCount() then
      startPtr = i
    end
  end

  local toRemove = {}
  while startPtr < endPtr do
    local from = entries[endPtr]
    local to = entries[startPtr]

    local emptied, _ = from.stack:transferTo(to.stack)
    if emptied then
      table.insert(toRemove, from.index)
      endPtr = endPtr - 1
    end

    if (not emptied) or (to.stack:getCount() == to.stack:getMaxCount()) then
        startPtr = startPtr + 1
    end
  end

  table.sort(toRemove, function(a, b) return a > b end)
  for _,index in ipairs(toRemove) do
    self:_removeStack(index)
  end

  return #toRemove
end

return ItemGroup