-- Constants and magic numbers
local mapgen_buffer = 16

-- Convert 3d relative coordinates to an index on a flat array
local function from_3d_to_flat(dx, dy, dz, nx, ny)
    return (nx * ny * dz) + (nx * dy) + dx + 1
end

-- Convert an index on a flat array to 3d relative coordinates
local function from_flat_to_3d(i, nx, ny)
    return {
        dx = (i - 1) % nx,
        dy = math.floor((i - 1) / nx) % ny,
        dz = math.floor((i - 1) / (nx * ny))
    }
end

-- Iterate over a 3d area, both indexing by the index and the ansolute position
local function iter_3d_area(minp, maxp, callback)
    local nx = maxp.x - minp.x + 1
    local ny = maxp.y - minp.y + 1
    local nz = maxp.z - minp.z + 1

    for i = 1, nx * ny * nz do
        local dpos = from_flat_to_3d(i, nx, ny)
        local pos = {
            x = minp.x + dpos.dx,
            y = minp.y + dpos.dy,
            z = minp.z + dpos.dz,
        }

        callback(i, pos)
    end
end

-- Helper function to convert a set of coordinates to a readable string
local function pos_to_str(pos)
    return "(" .. pos.x .. ", " .. pos.y .. ", " .. pos.z .. " )"
end

local Flat3dArray = {}
Flat3dArray.__index = Flat3dArray

function Flat3dArray:new(minp, maxp, arr)
    local instance = {}
    setmetatable(instance, Flat3dArray)

    local nx = maxp.x - minp.x + 1
    local ny = maxp.y - minp.y + 1
    local nz = maxp.z - minp.z + 1

    if #arr ~= nx * ny * nz then
        error(
            "Input array doesn't match dimension lengths: " .. nx .. " x " ..
            ny .. " x " .. nz .. " = " .. (nx*ny*nz) .. ", but found " .. #arr
        )
    end

    instance.nx = nx
    instance.ny = ny
    instance.nz = nz
    instance.minp = minp
    instance.maxp = maxp
    instance.arr = arr

    return instance
end

function Flat3dArray:from_func(minp, maxp, callback)
    local arr = {}

    iter_3d_area(minp, maxp, function (i, pos)
        arr[i] = callback(i, pos)
    end)

    return self:new(minp, maxp, arr)
end

function Flat3dArray:get_index(i)
    local out = self.arr[i]
    if out == nil then
        error(
            "Index " .. i .. " not found in array of length " .. #self.arr
        )
    end
    return out
end

function Flat3dArray:get_pos(pos)
    self:validate_pos(pos)

    local dx = pos.x - self.minp.x
    local dy = pos.y - self.minp.y
    local dz = pos.z - self.minp.z

    return self:get_index(from_3d_to_flat(dx, dy, dz, self.nx, self.ny))
end

local function is_valid_pos(pos, minp, maxp)
    if pos.x < minp.x then
        return false
    elseif pos.x > maxp.x then
        return false
    elseif pos.y < minp.y then
        return false
    elseif pos.y > maxp.y then
        return false
    elseif pos.z < minp.z then
        return false
    elseif pos.z > maxp.z then
        return false
    else
        return true
    end
end

function Flat3dArray:valid_pos(pos)
    return is_valid_pos(pos, self.minp, self.maxp)
end

function Flat3dArray:validate_pos(pos)
    if not self:valid_pos(pos) then
        error(
            table.concat(
                { "Position "
                , pos_to_str(pos)
                , " out of bounds from minp = "
                , pos_to_str(self.minp)
                , ", maxp = "
                , pos_to_str(self.maxp)
                },
                ""
            )
        )
    end
end

-- Get an enhanced function from the def function that warns us when the
-- function does not behave properly
local function enhanced_func(def)
    if type(def.name) ~= "string" then
        error("Invalid nameless shape definition")
    elseif type(def.func) ~= "function" then
        error(
            "Invalid shape definition misses an adjustment function"
        )
    else
        return function(pos, n)
            local out = def.func(pos, n)

            if type(out) == "number" then
                return out
            elseif n == nil then
                error(
                    "Shape " .. def.name .. " function must return a number. The input `n` was nil. Perhaps your `noise_params` field is invalid?"
                )
            else
                error("Shape " .. def.name .. " function must return a number.")
            end
        end
    end
end

-- Get a flat array of cave shape noise values from a given cave shape def
local function get_flat_from_shape_def(def, minp, maxp)
    local f = enhanced_func(def)

    local nx = maxp.x - minp.x + 1
    local ny = maxp.y - minp.y + 1
    local nz = maxp.z - minp.z + 1

    local noise_flat_map = {}

    -- If noise parameters have been defined, fill the table with noise
    -- If not, all values remain nil
    if def.noise_params ~= nil then
        local p = PerlinNoiseMap(def.noise_params, { x = nx, y = ny, z = nz })
        
        if nz == 1 then
            p:get_2d_map_flat(minp, noise_flat_map)
        else
            p:get_3d_map_flat(minp, noise_flat_map)
        end
    end
    
    iter_3d_area(minp, maxp, function(i, pos)
        noise_flat_map[i] = f(pos, noise_flat_map[i])
    end)

    return Flat3dArray:new(minp, maxp, noise_flat_map)
end

local function get_flat_from_noise_params(minp, maxp, noise_params)
    local nx = maxp.x - minp.x + 1
    local ny = maxp.y - minp.y + 1
    local nz = maxp.z - minp.z + 1

    local buffer = {}

    local p = PerlinNoiseMap(noise_params, { x = nx, y = ny, z = nz })

    if nz == 1 then
        p:get_2d_map_flat(minp, buffer)
    else
        p:get_3d_map_flat(minp, buffer)
    end

    return Flat3dArray:new(minp, maxp, buffer)
end

-- Based on the number of cave shapes, calculate how quickly connectivity is
-- meant to change
local function get_connectivity_noise_params(shape_size)
    local factor = math.max(math.abs(shape_size) ^ 0.5, 1)

    return {
        offset = 50,
        scale = 50,
        spread = { x = factor * 250, y = factor * 100, z = factor * 250 },
        seed = 297948,
        octaves = 2,
        persistence = 0.2,
        lacunarity = 2.0,
        flags = "eased"
    }
end

-- Based on the number of cave shapes, calculate how quickly verticality is
-- meant to change
local function get_verticality_noise_params(shape_size)
    local factor = math.max(math.abs(shape_size) ^ 0.5, 1)

    return {
        offset = 50,
        scale = 50,
        spread = { x = factor * 100, y = factor * 250, z = factor * 100 },
        seed = 35644,
        octaves = 2,
        persistence = 0.2,
        lacunarity = 2.0,
        flags = "eased"
    }
end

-- Get the distance of each cave shape to calculate their weight.
local function shape_distance_weight(shape_size, dx, dy)
    local factor = math.max(math.abs(shape_size) ^ 0.5, 1)
    local max_distance = 100 / factor

    if dx^2 + dy^2 > max_distance^2 then
        return 0
    else
        return 1 / (dx^5 + dy^5)
    end
end

-- Get a flat map containing all flat threshold values
local function get_threshold_flat(minp, maxp)
    local connectivity = get_flat_from_noise_params(
        minp, maxp,
        get_connectivity_noise_params(#noordstar_caves.registered_shapes)
    )
    local verticality = get_flat_from_noise_params(
        minp, maxp,
        get_verticality_noise_params(#noordstar_caves.registered_shapes)
    )

    local noise = {}

    -- Get noise for all cave shapes
    for key, def in pairs(noordstar_caves.registered_shapes) do
        noise[key] = {
            key = key,
            noise = get_flat_from_shape_def(def, minp, maxp),
            def = def
        }
    end

    -- Create the flat array
    return Flat3dArray:from_func(minp, maxp, function(i, pos)
        local total = 0
        local count = 0

        local x = connectivity:get_pos(pos)
        local y = verticality:get_pos(pos)

        for _, n in pairs(noise) do
            local v = n.noise:get_pos(pos)

            local dx = math.abs(x - n.def.connectivity_point)
            local dy = math.abs(y - n.def.verticality_point)

            local w = math.abs(shape_distance_weight(
                #noordstar_caves.registered_shapes, dx, dy
            ))

            total = total + v * w
            count = count + w
        end

        if count <= 0 then
            return -1000
        else
            return total / count
        end
    end)
end


-- Old minimum height for the overworld
local old_overworld_min = mcl_vars.mg_overworld_min

-- If another mod doesn't override the maximum world depth, we will assume that
-- the world depth is the following value.
local world_depth = -60

-- Otherwise, this variable can be changed using the following function,
-- which will also update all the other necessary variables
function noordstar_caves.set_world_depth(h)
    -- Set world depth variable
    world_depth = h
end
noordstar_caves.set_world_depth(world_depth)

-- The `cave_vastness` is a variable that determines the size of caves.
-- The function takes in an x, y and z variable, and returns a number between
-- 0 and 1.
-- Low numbers mean very rare and tiny caves, while high numbers mean massive
-- caves, with massive open spaces.
-- If you wish to overwrite this function, it is good to keep in mind to:
--  - Make sure that the output changes VERY SLOWLY over time
--  - This function will be run a LOT so it is very performance sensitive
noordstar_caves.cave_vastness = function(pos)
    return pos.y / world_depth
end

-- Secretly, we're using an internal function that also adds a safe layer for
-- when we're approaching bedrock levels
local function cave_vastness(pos)
    if world_depth + 20 < pos.y then
        return noordstar_caves.cave_vastness(pos)
    elseif world_depth + 5 < pos.y then
        return noordstar_caves.cave_vastness(pos) * math.abs(pos.y - world_depth - 5) / 15
    else
        return 0
    end
end

-- Get a flat array of nodes that are either in caves or not in caves
local function get_flat_cave_bools(minp, maxp)
    local thresholds = get_threshold_flat(minp, maxp)

    return Flat3dArray:from_func(minp, maxp, function(i, pos)
        return thresholds:get_pos(pos) >= 1 - cave_vastness(pos)
    end)
end

local node_type =
    { unknown = 1  -- Edge of a chunk
    , floor = 2    -- Floor of a cave
    , wall = 3     -- Side wall of a cave
    , roof = 4     -- Roof of a cave
    , content = 5  -- Air in the cave not adjacent to a wall, floor or roof
    , stone = 6    -- Underground node not adjacent to a cave
    }

local function get_flat_cave_node_types(minp, maxp)
    local bools = get_flat_cave_bools(minp, maxp)

    return Flat3dArray:from_func(minp, maxp, function (i, pos)
        if not bools:get_pos(pos) then
            return node_type.stone
        -- Floor takes precedence
        elseif pos.y == minp.y then -- Could be floor
            return node_type.unknown
        elseif not bools:get_pos({ x = pos.x, y = pos.y - 1, z = pos.z }) then
            return node_type.floor
        -- Then roof takes precedence
        elseif pos.y == maxp.y then -- Could be roof
            return node_type.unknown
        elseif not bools:get_pos({ x = pos.x, y = pos.y + 1, z = pos.z }) then
            return node_type.roof
        else
            -- Check for walls
            local left  = { x = pos.x - 1, y = pos.y, z = pos.z }
            local right = { x = pos.x + 1, y = pos.y, z = pos.z }
            local front = { x = pos.x, y = pos.y, z = pos.z - 1 }
            local back  = { x = pos.x, y = pos.y, z = pos.z + 1 }

            -- Check if the value is near the edge
            local on_edge = false

            if left.x < minp.x then
                on_edge = true
            elseif not bools:get_pos(left) then
                return node_type.wall
            end
            if right.x > maxp.x then
                on_edge = true
            elseif not bools:get_pos(right) then
                return node_type.wall
            end
            if front.z < minp.z then
                on_edge = true
            elseif not bools:get_pos(front) then
                return node_type.wall
            end
            if back.z > maxp.z then
                on_edge = true
            elseif not bools:get_pos(back) then
                return node_type.wall
            end

            if on_edge then
                return node_type.unknown
            else
                return node_type.content
            end
        end
    end)
end

minetest.register_on_generated(function(minp, maxp, blockseed)
    local vminp =
        { x = minp.x - mapgen_buffer
        , y = minp.y - mapgen_buffer
        , z = minp.z - mapgen_buffer
        }
    local vmaxp =
        { x = maxp.x + mapgen_buffer
        , y = maxp.y + mapgen_buffer
        , z = maxp.z + mapgen_buffer
        }

    -- Get voxelmanip
    local vm = minetest.get_mapgen_object("voxelmanip")
    local flat_data = Flat3dArray:new(vminp, vmaxp, vm:get_data())
    
    -- Get threshold values
    local node_types = get_flat_cave_node_types(vminp, vmaxp)
    
    local node_air   = minetest.get_content_id("air")
    local node_floor = minetest.get_content_id("mcl_core:glass_green")
    local node_wall  = minetest.get_content_id("mcl_core:glass_purple")
    local node_roof  = minetest.get_content_id("mcl_core:glass_red")
    local node_other = minetest.get_content_id("mcl_core:glass")

    local nids = Flat3dArray:from_func(vminp, vmaxp, function(i, pos)
        local nt = node_types:get_pos(pos)

        if not is_valid_pos(pos, minp, maxp) then
            return flat_data:get_pos(pos)
        elseif nt == node_type.unknown then
            return node_other
        elseif nt == node_type.floor then
            return node_floor
        elseif nt == node_type.wall then
            return node_wall
        elseif nt == node_type.roof then
            return node_roof
        elseif nt == node_type.content then
            return node_air
        elseif nt == node_type.stone then
            return flat_data:get_pos(pos)
        end
    end)

    -- Write all changes to the Minetest world
    vm:set_data(nids.arr)
    vm:write_to_map()
end)