require "ISUI/ISInventoryPaneContextMenu"

local genuine_ISInventoryPaneContextMenu_doMagazineMenu = ISInventoryPaneContextMenu.doMagazineMenu;
ISInventoryPaneContextMenu.doMagazineMenu = function(playerObj, magazine, context)
    genuine_ISInventoryPaneContextMenu_doMagazineMenu(playerObj, magazine, context)

	context:addOption(getText("ContextMenu_ReloadAllMagazines_LoadAllMagazines", magazine:getName()), playerObj, ISInventoryPaneContextMenu.onLoadBulletsInAllMagazines, magazine);
	context:addOption(getText("ContextMenu_ReloadAllMagazines_UnloadAllMagazines", magazine:getName()), playerObj, ISInventoryPaneContextMenu.onUnloadBulletsFromAllMagazines, magazine);
end

ISInventoryPaneContextMenu.onLoadBulletsInAllMagazines = function(playerObj, magazine)
	local allAmmoCount = playerObj:getInventory():getItemCountRecurse(magazine:getAmmoType());
	local magazineCount = playerObj:getInventory():getItemCountRecurse(magazine:getType());
	local maxAmmo = magazine:getMaxAmmo();
	
	local ammoToGet = magazineCount * maxAmmo;
	if magazineCount * maxAmmo > allAmmoCount then
		ammoToGet = allAmmoCount;
	end

	--local items = playerObj:getInventory():getSomeTypeRecurse(magazine:getAmmoType(), ammoToGet)
	--ISInventoryPaneContextMenu.transferIfNeeded(playerObj, items)
	local allAmmoItems = playerObj:getInventory():getSomeTypeRecurse(magazine:getAmmoType(), allAmmoCount)


	local ammoIndex = 0;
	local containers = {};
    for i = 0, playerObj:getInventory():getItems():size() -1 do
        local item = playerObj:getInventory():getItems():get(i);
		if instanceof(item, "InventoryContainer") then
			table.insert(containers, item)
		end
        if item:getType() == magazine:getType() then
            local magazine = item;
            local ammoCount = playerObj:getInventory():getItemCountRecurse(magazine:getAmmoType());
            if ammoCount > magazine:getMaxAmmo() then
                ammoCount = magazine:getMaxAmmo();
            end
            if ammoCount > magazine:getMaxAmmo() - magazine:getCurrentAmmoCount() then
                ammoCount = magazine:getMaxAmmo() - magazine:getCurrentAmmoCount();
            end
            if ammoCount ~= 0 then
				ISInventoryPaneContextMenu.transferIfNeeded(playerObj, magazine)
				local ammoToLoad = ISInventoryPaneContextMenu.getUniqueBullets(allAmmoItems, allAmmoCount, ammoIndex, ammoCount)
				ammoIndex = ammoIndex + ammoCount
				ISInventoryPaneContextMenu.transferIfNeeded(playerObj, ammoToLoad)
				ISTimedActionQueue.add(ISLoadBulletsInMagazine:new(playerObj, magazine, ammoCount))
            end
        end
    end
	
    for k, v in pairs(containers) do
		local container = v:getInventory():getItems();
		for i = 0, container:size() -1 do
			local item = container:get(i);
			if item:getType() == magazine:getType() then
				local magazine = item;
				local ammoCount = playerObj:getInventory():getItemCountRecurse(magazine:getAmmoType());
				if ammoCount > magazine:getMaxAmmo() then
					ammoCount = magazine:getMaxAmmo();
				end
				if ammoCount > magazine:getMaxAmmo() - magazine:getCurrentAmmoCount() then
					ammoCount = magazine:getMaxAmmo() - magazine:getCurrentAmmoCount();
				end
				if ammoCount ~= 0 then
					ISInventoryPaneContextMenu.transferIfNeeded(playerObj, magazine)
					local ammoToLoad = ISInventoryPaneContextMenu.getUniqueBullets(allAmmoItems, allAmmoCount, ammoIndex, ammoCount)
					ammoIndex = ammoIndex + ammoCount
					ISInventoryPaneContextMenu.transferIfNeeded(playerObj, ammoToLoad)
					ISTimedActionQueue.add(ISLoadBulletsInMagazine:new(playerObj, magazine, ammoCount))
				end
			end
		end
    end
end

ISInventoryPaneContextMenu.getUniqueBullets = function(allAmmoItems, allAmmoCount, ammoIndex, ammoCount)
	local ammoToLoad = ArrayList.new()
	for i = ammoIndex, ammoIndex + ammoCount - 1 do
		if i >= allAmmoCount then return ammoToLoad end
		local item = allAmmoItems:get(i);
		ammoToLoad:add(item);
	end
	return ammoToLoad
end

ISInventoryPaneContextMenu.onUnloadBulletsFromAllMagazines = function(playerObj, magazine)
	local containers = {};
    for i = 0, playerObj:getInventory():getItems():size() -1 do
        local item = playerObj:getInventory():getItems():get(i);
		if instanceof(item, "InventoryContainer") then
			table.insert(containers, item)
		end
        if item:getType() == magazine:getType() then
            local magazine = item;
			if magazine:getCurrentAmmoCount() > 0 then
				ISInventoryPaneContextMenu.transferIfNeeded(playerObj, magazine)
				ISTimedActionQueue.add(ISUnloadBulletsFromMagazine:new(playerObj, magazine))
			end
        end
    end
	
    for k, v in pairs(containers) do
		local container = v:getInventory():getItems();
		for i = 0, container:size() -1 do
			local item = container:get(i);
			if item:getType() == magazine:getType() then
				local magazine = item;
				if magazine:getCurrentAmmoCount() > 0 then
					ISInventoryPaneContextMenu.transferIfNeeded(playerObj, magazine)
					ISTimedActionQueue.add(ISUnloadBulletsFromMagazine:new(playerObj, magazine))
				end
			end
		end
    end
end