-- Copyright (C) 2025  snoutie
-- Authors: snoutie (copyright@achtarmig.org)
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as published
-- by the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.

-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.

-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <https://www.gnu.org/licenses/>.

local modpath = core.get_modpath(core.get_current_modname())
local internal = dofile(modpath .. "/internal.lua")

-- A table of offsets for ease of use
local offsets = {}
-- local offsets = {
--     { x = 0,  z = 0,  y = 0, d = 0 },

--     { x = 0,  z = 1,  y = 0, d = 1 },
--     { x = 0,  z = -1, y = 0, d = 1 },

--     { x = 1,  z = 0,  y = 0, d = 1 },
--     { x = -1, z = 0,  y = 0, d = 1 },
-- }
for x = -1, 1 do
    for y = -1, 1 do
        for z = -1, 1 do
            local d = math.abs(x) + math.abs(y) + math.abs(z)
            table.insert(offsets, {
                x = x, -- X Offset
                y = y, -- Y Offset
                z = z, -- Z Offset
                d = d  -- Distance to 0 / 0 / 0
            })
        end
    end
end

-- Store number of offsets for less calculations
local number_of_offsets = table.getn(offsets)

-- Adds offset to pos
-- @param pos    table Position
-- @param offset table Offset
-- @return       table Position with added offset
local function add_offset_to_pos(pos, offset)
    return {
        x = pos.x + offset.x,
        y = pos.y + offset.y,
        z = pos.z + offset.z,
    }
end

-- Get a LPN from either the buffer or create a new one
--  and add it to the buffer
-- @param b   table Buffered LPN
-- @param pos table Positon of the node
-- @return    table LPN
local function get_lpn_buffered(b, pos)
    local h = core.hash_node_position(pos)
    if b[h] ~= nil then
        return b[h]
    end
    b[h] = internal.new_lpn(h, pos)
    return b[h]
end

local function round_away_from_zero(value)
    if value == 0 then
        return 0
    end
    if value > 0 then
        return math.ceil(value)
    end
    return math.floor(value) -- Let's hope math floor is optimized for zero
end

-- Checks if liquid_id of curr and lpn are the same.
-- Checks if the liquid_level is above -1
-- @param curr table The current LPN
-- @param lpn  table The LPN to compare against
-- @return     bool  Is this LPN relevant
local function is_lpn_relevant(curr, lpn)
    return lpn.liquid_level >= 0 and (lpn.liquid_level == 0 or curr.liquid_id == lpn.liquid_id)
end

-- local function get_gradient(base, compare, opp)
--     local compare_copy = compare
--     if compare_copy < 0 then compare_copy = 0 end
--     return math.ceil(((opp + base) - compare_copy * 2) / 8)
-- end

-- Gets all adjacent liquid nodes with the same liquid_id as curr
-- Gets all adjacent air nodes
-- @param b     table  Buffered LPN
-- @param curr  table  The current LPN
-- @return      table  Table of neighboring LPN
-- @return      number Number of neighbors
local function get_all_neighbors_pressure(b, curr)
    local neighbors = {}
    local no_of_h_neighbors = 0
    local no_of_v_neighbors = 0
    for i = 1, number_of_offsets do
        local current_offset = offsets[i]
        local pos = add_offset_to_pos(curr.pos, current_offset)
        local lpn = get_lpn_buffered(b, pos)

        --if is_lpn_relevant(curr, lpn) then

        local o_off = {}
        if i % 2 == 0 then
            o_off = offsets[i - 1]
        else
            o_off = offsets[i + 1]
        end

        local o_pos = add_offset_to_pos(curr.pos, o_off)
        local o_lpn = get_lpn_buffered(b, o_pos)

        local g = current_offset.g + get_gradient(curr.liquid_level, lpn.liquid_level, o_lpn.liquid_level)

        if g ~= 0 then
            if i < 5 then
                no_of_h_neighbors = no_of_h_neighbors + 1
            else
                no_of_v_neighbors = no_of_v_neighbors + 1
            end
        end
        table.insert(neighbors, { lpn = lpn, g = g })
        --end
    end
    return neighbors, no_of_h_neighbors, no_of_v_neighbors
end


local function get_pressure_gradient(b, lpn)
    local number_of_directions = 0

    lpn.neighbors = get_all_neighbors_pressure(b, lpn)

    lpn.pressure_gradient = internal.get_inital_pressure_gradient()

    local i = 1
    for _, g in pairs(lpn.pressure_gradient) do
        g = g + get_gradient(lpn.liquid_level, lpn.neighbors[i].liquid_level)
        i = i + 1
        if g ~= 0 then
            number_of_directions = number_of_directions + 1
        end
    end

    return number_of_directions
end

local function get_density_influence(r, d)
    local vol = math.pi * (r ^ 8) / 4
    local v = math.max(0, r - d)
    return v + v * v / vol
end

local function get_density_influence_derivative(r, d)
    if (d >= r) then return 0 end
    local f = r * r - d * d
    local scale = -24 / (math.pi * (r ^ 8))
    return scale * d * f * f
end

local function get_density(b, lpn)
    local radius = 3 -- 3 is max
    for i = 1, number_of_offsets do
        local offset = offsets[i]
        local o_pos = add_offset_to_pos(lpn.pos, offset)
        local o_lpn = get_lpn_buffered(b, o_pos)

        local liquid_level = o_lpn.liquid_level
        if liquid_level < 0 then
            liquid_level = lpn.liquid_level
        elseif liquid_level == 0 then
            liquid_level = 0
        end

        -- if offset.d == 0 then
        --     liquid_level = 0
        -- end

        local influence = liquid_level * get_density_influence(radius, offset.d)

        table.insert(lpn.neighbors, { lpn = o_lpn, prs = 0 })

        lpn.density = lpn.density + influence
    end
    lpn.density = math.floor(lpn.density)
end


local function get_pressure_from_density(density)
    local multiplier = 1
    local target_pressure = 11
    --core.chat_send_all(density)
    return (density - target_pressure) * multiplier
end

local function get_gradient(b, lpn)
    -- local gradient = {
    --     x = 0,
    --     y = 0,
    --     z = 0,
    -- }
    local radius = 3 -- 3 is max

    for i = 1, number_of_offsets do
        local offset = offsets[i]

        if offset.d == 0 then
            goto continue
        end

        local neighbor = lpn.neighbors[i].lpn

        local density = neighbor.density

        -- if density < 0 then
        --     density = neighbor.density
        -- else
        if density == 0 then
            density = 1
        end

        local slope = get_density_influence_derivative(radius, offset.d)
        local pressure = get_pressure_from_density(density)

        --if offset.y == 0 then
        lpn.neighbors[i].prs = math.ceil(math.abs(
            (pressure * slope * lpn.liquid_level / density)))

        lpn.neighbors[i].is_down = false

        if offset.y == 1 then
            lpn.neighbors[i].prs = lpn.neighbors[i].prs - 6
        end

        if offset.y == -1 then
            lpn.neighbors[i].prs = 1
            lpn.neighbors[i].is_down = true
        end

        -- lpn.gradient = {
        --     x = lpn.gradient.x + (pressure * offset.x * slope * neighbor.liquid_level / density),
        --     y = lpn.gradient.y + (pressure * offset.y * slope * neighbor.liquid_level / density),
        --     z = lpn.gradient.z + (pressure * offset.z * slope * neighbor.liquid_level / density),
        -- }
        ::continue::
    end

    --core.chat_send_all("Pressure: " .. dump(lpn.gradient))

    -- lpn.gradient = {
    --     x = round_away_from_zero(lpn.gradient.x),
    --     y = round_away_from_zero(lpn.gradient.y),
    --     z = round_away_from_zero(lpn.gradient.z),
    -- }
    -- if lpn.gradient.x > 0 then
    --     lpn.gradient.x = 1
    -- elseif lpn.gradient.x < 0 then
    --     lpn.gradient.x = -1
    -- end

    -- if lpn.gradient.y > 0 then
    --     lpn.gradient.y = 1
    -- elseif lpn.gradient.y < 0 then
    --     lpn.gradient.y = -1
    -- end

    -- if lpn.gradient.z > 0 then
    --     lpn.gradient.z = 1
    -- elseif lpn.gradient.z < 0 then
    --     lpn.gradient.z = -1
    -- end
end


local function get_pressure_force(lpn)

end

-- Gets neighbors of curr and calculates the pressure
-- This function combines get_valid_neighbors and get_pressure_straight
-- @param b    table Buffered LPN
-- @param curr table The current LPN
-- @return     float Pressure of curr calculated from valid neighbors
local function get_valid_neigbor_pressure(b, curr)
    local number_of_neighbors = 0
    local total_liquid_level = curr.liquid_level
    for i = 1, number_of_offsets do
        local pos = add_offset_to_pos(curr.pos, offsets[i])
        local lpn = get_lpn_buffered(b, pos)

        if is_lpn_relevant(curr, lpn) then
            total_liquid_level = total_liquid_level + lpn.liquid_level / 2
            number_of_neighbors = number_of_neighbors + 1
        end
    end
    return total_liquid_level / (number_of_neighbors + 1)
end

-- Gets the pressure of curr from the neighbors given
-- @param b                   table  Buffered LPN
-- @param neighbors           table  Table of neighboring LPN
-- @param number_of_neighbors number Number of neighbors
-- @param curr                table  The current LPN
-- @return                    float  Pressure of curr calculated from neighbors
local function get_pressure_straight(b, neighbors, number_of_neighbors, curr)
    local total_liquid_level = curr.liquid_level

    for i = 1, number_of_neighbors do
        total_liquid_level = total_liquid_level + neighbors[i].liquid_level / 2
    end
    return total_liquid_level / (number_of_neighbors + 1)
end

-- Tries to move amount from "from" to "to"
-- Actual amount moved may not be the same as "amount"
--  therefore the actual amount moved is returned
-- @param from table  LPN from which amount is deducted
-- @param to   table  LPN to which amount is added
-- @return     number Actual amount moved
local function try_move(from, to, amount)
    if to.liquid_level >= 8 then
        return 0
    end
    if amount <= 0 then
        return 0
    end
    local max_allowed = amount + 8 - (to.liquid_level + amount)
    local max_allowed_clamped = math.min(max_allowed, amount, from.liquid_level)

    internal.set_lpn(from.liquid_id, from, from.liquid_level - max_allowed_clamped)
    internal.set_lpn(from.liquid_id, to, to.liquid_level + max_allowed_clamped)

    if max_allowed_clamped ~= amount then
        return max_allowed_clamped
    end
    return amount
end

-- Calculates how to move curr_pos
-- @param b        table Buffered LPN
-- @param curr_pos table Current position
local function move(b, curr_pos)
    local curr_nbs
    local curr_nnbs
    local up_lpn

    --check down first
    local down_lpn = get_lpn_buffered(b, { x = curr_pos.x, y = curr_pos.y - 1, z = curr_pos.z })
    local curr_lpn = get_lpn_buffered(b, curr_pos)

    if curr_lpn.liquid_level <= 0 then
        internal.remove_node_to_check(curr_lpn.pos)
        return
    end

    curr_nbs, curr_nnbs = get_valid_neighbors(b, curr_lpn)

    if is_lpn_relevant(curr_lpn, down_lpn) then
        local moved = try_move(curr_lpn, down_lpn, curr_lpn.liquid_level)
        if moved > 0 then
            up_lpn = get_lpn_buffered(b, { x = curr_pos.x, y = curr_pos.y + 1, z = curr_pos.z })
            if up_lpn.liquid_level >= 0 then
                internal.add_node_to_check(up_lpn.pos)
            end
            internal.add_node_to_check(down_lpn.pos)
        end
        if moved == curr_lpn.liquid_level then
            internal.remove_node_to_check(curr_lpn.pos)
            for i = 1, curr_nnbs do
                internal.add_node_to_check(curr_nbs[i].pos)
            end
            return
        end
    end

    -- Every neighbor will be higher (except for air) so processing is done
    if curr_lpn.liquid_level <= 1 then
        internal.remove_node_to_check(curr_lpn.pos)
        return
    end

    local curr_prs = get_pressure_straight(b, curr_nbs, curr_nnbs, curr_lpn)

    if curr_prs >= 4.8 then --4.8 == ((8*4/2)+8)/5 -> meaning maximum pressure, implies no way to move
        internal.remove_node_to_check(curr_lpn.pos)
        return
    end

    -- TODO: Understand what is happening here and optimize
    local number_of_swaps = 0
    for i = 1, curr_nnbs do
        local next_lpn = curr_nbs[i]
        local next_prs = get_valid_neigbor_pressure(b, next_lpn)
        if curr_prs > next_prs and next_lpn.liquid_level < curr_lpn.liquid_level then
            if try_move(curr_lpn, next_lpn, 1) > 0 then
                internal.add_node_to_check(next_lpn.pos)
                internal.add_node_to_check(
                    get_lpn_buffered(b, { x = curr_pos.x, y = curr_pos.y + 1, z = curr_pos.z }).pos
                )
                number_of_swaps = number_of_swaps + 1
            end
        elseif curr_lpn.liquid_level < next_lpn.liquid_level then
            internal.add_node_to_check(next_lpn.pos)
        end
    end

    if number_of_swaps == 0 then
        internal.remove_node_to_check(curr_lpn.pos)
    end
end

local function try_spread(buffer, from, to, is_down)
    if is_lpn_relevant(from, to) and ((from.liquid_level > 1 and from.liquid_level >= to.liquid_level + 1) or is_down) then
        return try_move(from, to, 1)
    end
    return 0
end

local function shuffle(array)
    math.randomseed(os.time())                  -- Seed the random number generator
    for i = #array, 2, -1 do
        local j = math.random(i)                -- Pick a random index from 1 to i
        array[i], array[j] = array[j], array[i] -- Swap elements
    end
end

local last_finish = 0
local active_blocks = {}

core.register_on_liquid_transformed(function(pos_list, node_list)
    for hash, block_start in pairs(active_blocks) do
        if core.compare_block_status(block_start, "active") then
        else
            active_blocks[hash] = nil
            goto continue
        end

        --core.chat_send_all(dump(block_start) .. "----" .. modified_blocks_count)
        local buffer = {}
        local block_end = {
            x = block_start.x + 15,
            y = block_start.y + 15,
            z = block_start.z + 15,
        }
        local node_pos_list = core.find_nodes_in_area(block_start, block_end, { "group:liquid_physics" })

        local lpn_list = {}
        local number_of_lpn = 0
        for _, pos in pairs(node_pos_list) do
            local lpn = get_lpn_buffered(buffer, pos)
            get_density(buffer, lpn)

            table.insert(lpn_list, lpn)
            number_of_lpn = number_of_lpn + 1
        end

        for i = 1, number_of_lpn do
            local g_lpn = lpn_list[i]
            get_gradient(buffer, g_lpn)

            --g_lpn.gradient.y = 0

            --core.chat_send_all("Grad: " .. dump(g_lpn.gradient))

            -- local spread_to_pos = add_offset_to_pos(g_lpn.pos, g_lpn.gradient)
            -- local spread_to = get_lpn_buffered(buffer, spread_to_pos)

            -- try_spread(buffer, g_lpn, spread_to)



            -- local gs = 0
            -- for g = 1, number_of_offsets do
            --     if g_lpn.neighbors[g].prs > 0 then
            --         gs = gs + 1
            --     end
            -- end
            local number_of_targets = 0
            for g = 1, number_of_offsets do
                if g_lpn.neighbors[g].prs > 0 and not g_lpn.neighbors[g].is_down then
                    number_of_targets = number_of_targets + 1
                end
            end

            if number_of_targets >= g_lpn.liquid_level then
                shuffle(g_lpn.neighbors)
            end

            for g = 1, number_of_offsets do
                if g_lpn.neighbors[g].prs > 0 then
                    try_spread(buffer, g_lpn, g_lpn.neighbors[g].lpn, g_lpn.neighbors[g].is_down)
                end
            end


            -- if g_lpn.gradient.x ~= 0 then
            --     local spread_to_pos = add_offset_to_pos(g_lpn.pos, { x = g_lpn.gradient.x, y = 0, z = 0 })
            --     local spread_to = get_lpn_buffered(buffer, spread_to_pos)

            --     try_spread(buffer, g_lpn, spread_to)
            --     -- if try_spread(buffer, g_lpn, spread_to, gs) == 0 then
            --     --     spread_to_pos = add_offset_to_pos(g_lpn.pos, { x = g_lpn.gradient.x * -1, y = 0, z = 0 })
            --     --     spread_to = get_lpn_buffered(buffer, spread_to_pos)
            --     --     try_spread(buffer, g_lpn, spread_to, gs)
            --     -- end
            -- end
            -- if g_lpn.gradient.z ~= 0 then
            --     local spread_to_pos = add_offset_to_pos(g_lpn.pos, { x = 0, y = 0, z = g_lpn.gradient.z })
            --     local spread_to = get_lpn_buffered(buffer, spread_to_pos)

            --     try_spread(buffer, g_lpn, spread_to)
            --     -- if try_spread(buffer, g_lpn, spread_to, gs) == 0 then
            --     --     spread_to_pos = add_offset_to_pos(g_lpn.pos, { x = 0, y = 0, z = g_lpn.gradient.z * -1 })
            --     --     spread_to = get_lpn_buffered(buffer, spread_to_pos)
            --     --     try_spread(buffer, g_lpn, spread_to, gs)
            --     -- end
            -- end
        end

        if lpn_list[1] ~= nil then
        end

        -- for _, pos in pairs(node_pos_list) do
        --     local lpn = get_lpn_buffered(buffer, pos)

        --     local neighbors, h, v = get_all_neighbors_pressure(buffer, lpn)

        --     local n = 0
        --     for i = 1, number_of_offsets do
        --         if neighbors[i].g == 0 then
        --             goto continue
        --         end
        --         if i < 5 then
        --             n = h
        --         else
        --             n = v - 1
        --         end
        --         try_spread(buffer, lpn, neighbors[i].lpn, n, neighbors[i].g)
        --         ::continue::
        --     end
        -- end

        -- ::continue::
        --
        for _, lpn in pairs(buffer) do
            if lpn.init_liquid_level ~= lpn.liquid_level then
                if lpn.liquid_id then
                    internal.set_node(lpn)
                end
            end
        end
        ::continue::
    end
end)

core.register_on_mapblocks_changed(function(modified_blocks, modified_blocks_count)
    for hash, _ in pairs(modified_blocks) do
        local block_start = core.get_position_from_hash(hash)

        block_start = {
            x = block_start.x * 16,
            y = block_start.y * 16,
            z = block_start.z * 16,
        }

        if core.compare_block_status(block_start, "active") then
            active_blocks[hash] = block_start
        else
            active_blocks[hash] = nil
        end
    end
    -- -- Buffered LPN
    -- local b = {}
    -- for hpos, pos in pairs(liquid_physics._nodes_to_check) do
    --     if core.compare_block_status(pos, "active") then
    --         move(b, pos)
    --     else
    --         internal.remove_node_to_check(pos)
    --     end
    -- end
    -- -- Update only changed nodes from buffer
    -- for _, lpn in pairs(b) do
    --     if lpn.init_liquid_level ~= lpn.liquid_level then
    --         internal.set_node(lpn)
    --     end
    -- end
end)