-- Basically: https://stackoverflow.com/a/1761646
local function weighted_random(tbl)
    local sum = 0
    for _, item in pairs(tbl) do
        sum = sum + item.weight
    end

    local r = math.random(0, sum - 1)
    for _, item in pairs(tbl) do
        if r < item.weight then
            return item
        end
        r = r - item.weight
    end

    assert(false, "shouldn't happen")
end

-- Returns a list of {name = "default:item", count = {min, max}} tables
-- It is up to the caller to randomize the quantity of each item
-- If 'include_empty' is true, also returns {skip = true} placeholders for
-- slots with nothing in them; callers should check for the existence of this
-- field and skip over the slot if it does exist or handle it in their own way
function yams.get_loot(pos_y, chest_type, include_empty)
    assert(yams.loot[chest_type], "[yams_mapgen] chest_type is bad")

    local possible_loot = {}
    for _, item in pairs(yams.loot[chest_type].items) do
        if (not item.y) or (pos_y >= item.y[1] and pos_y <= item.y[2]) then
            table.insert(possible_loot, item)
        end
    end

    local res = {}
    for i = 1, yams.loot[chest_type].slots do
        if math.random() < yams.loot[chest_type].slot_chance then
            -- Copy it so that the client cannot change it on accident
            local pick = table.copy(weighted_random(possible_loot))
            table.insert(res, {name = pick.name, count = pick.count})
        elseif include_empty then
            table.insert(res, {skip = true})
        end
    end

    return res
end

function yams.override_loot_funcs()
    if dungeon_loot._internal_get_loot then
        -- dungeontype is unused in this impl
        dungeon_loot._internal_get_loot = function(pos_y, dungeontype)
            local res = yams.get_loot(pos_y, "dungeon")
            for _, item in pairs(res) do
                -- Force the code to accept the item since we already
                -- determined the items randomly
                item.chance = 1.0
            end
            return res
        end
    end

    if tsm_pyramids.fill_chest then
        -- treasure_chance parameter is ignored
        tsm_pyramids.fill_chest = function(pos, stype, flood_sand, _)
            local n = core.get_node(pos)
            if not n or not n.name or n.name ~= "default:chest" then
                return false
            end

            local inv = core.get_meta(pos):get_inventory()
            inv:set_size("main", 8 * 4)

            local loot = yams.get_loot(pos.y, "tsm_pyramids", true)
            for i = 1, #loot do
                local item = loot[i]

                -- If the pyramid is flooded, fill every empty slot with sand
                if item.skip and flood_sand then
                    local sand = "default:sand"
                    if stype == "desert_sandstone" or stype == "desert_stone" then
                        sand = "default:desert_sand"
                    end
                    inv:set_stack("main", i, sand)
                elseif not item.skip then
                    local count = math.random(item.count[1], item.count[2])
                    local stack = {name = item.name, count = count}
                    inv:set_stack("main", i, stack)
                end
            end

            -- Used for awards
            local meta = core.get_meta(pos)
            meta:set_string("yams_loot_chest_type", "tsm_pyramids")

            return true
        end
    end

    if tsm_railcorridors then
        tsm_railcorridors.get_treasures = function(pr, pos)
            -- Include empty slots because otherwise the items will all be
            -- next to each other, and I prefer them to be scattered around
            -- the chest
            local loot = yams.get_loot(pos.y, "tsm_railcorridors", true)

            local res = {}
            for i = 1, #loot do
                local item = loot[i]
                if item.skip then
                    -- Empty slots are represented by the empty string
                    table.insert(res, "")
                else
                    local count = math.random(item.count[1], item.count[2])
                    local stack = {name = item.name, count = count}
                    table.insert(res, stack)
                end
            end

            -- Used for awards
            local meta = core.get_meta(pos)
            meta:set_string("yams_loot_chest_type", "tsm_railcorridors")

            return res
        end
    end
end

local function loot_check()
    local dump_item_probs = false

    -- caveat: depth restrictions are not considered
    if dump_item_probs then
        for name, data in pairs(yams.loot) do
            print(name)
            print("---")

            local sum = 0
            for _, item in pairs(data.items) do
                sum = sum + item.weight
            end

            for _, item in pairs(data.items) do
                local s = string.format("%.2f%%", (item.weight / sum * 100))
                print(item.name, s)
            end

            print()  -- new line
        end
    end

    -- Some basic checks
    for name, data in pairs(yams.loot) do
        for _, item in pairs(data.items) do
            if not core.registered_items[item.name] then
                assert(false, "[yams_mapgen] " .. item.name .. " is bad")
            end
            if item.count and item.count[2] < item.count[1] then
                assert(false, "[yams_mapgen] " .. item.name .. ": count values")
            end
            if item.y and item.y[2] < item.y[1] then
                assert(false, "[yams_mapgen] " .. item.name .. ": y values")
            end
        end
    end
end

core.register_on_mods_loaded(loot_check)
core.register_on_mods_loaded(yams.override_loot_funcs)
