-- by Sovereignty 12/11/2024

local random = Random.new()


export type Grid<e> = {
    Nodes: { {
        Entry: e,
        Index: number,
        RowX: number,
        RowZ: number,
    } },
    Index: (index: number) -> (GridMethods),
    FirstRow: (x: number) -> (number),
    LastRow: (x: number) -> (number),
}

type GridMethods = {
    GetNext: (self: {}) -> (number?),
    GetPrev: (self: {}) -> (number?),
    GetForward: (self: {}) -> (number?),
    GetBackward: (self: {}) -> (number?),
}

local function new<e>(rowX: number, rowZ: number, perGridEntryCallback: (index: number, x: number, z: number) -> (e)): Grid<e>
    local grid = {}

    for zIndex = 1, rowZ do
        for xIndex = 1, rowX do
            local thisPadIndex = (zIndex - 1) * rowX + xIndex
    
            table.insert(grid, {
                Entry = perGridEntryCallback(thisPadIndex, xIndex, zIndex),
                Index = thisPadIndex,
                RowX = xIndex,
                RowZ = zIndex
            })
        end
    end

    local entryIndex: number

    local function getNextEntry(): number?
        local entryData = grid[entryIndex]
        local xIndex: number = entryData.RowX
        local zIndex: number = entryData.RowZ
        if xIndex < rowX then
            return (zIndex - 1) * rowX + (xIndex + 1)
        end
        
        return nil
    end

    local function getPrevEntry(): number?
        local entryData = grid[entryIndex]
        if entryData.Index == entryIndex then
            local xIndex: number = entryData.RowX
            local zIndex: number = entryData.RowZ
            if xIndex > 1 then
                return (zIndex - 1) * rowX + (xIndex - 1)
            end
        end
            
        return nil
    end

    local function getForwardEntry(): number?
        local entryData = grid[entryIndex]
        if entryData.Index == entryIndex then
            local xIndex: number = entryData.RowX
            local zIndex: number = entryData.RowZ
            if zIndex < rowZ then
                return zIndex * rowX + xIndex
            end
        end
            
        return nil
    end

    local function getBackwardEntry(): number?
        local entryData = grid[entryIndex]
        if entryData.Index == entryIndex then
            local xIndex: number = entryData.RowX
            local zIndex: number = entryData.RowZ
            if zIndex > 1 then
                return (zIndex - 2) * rowX + xIndex
            end
        end

        return nil
    end

    local function firstRowX(entry: number): number
        return (1 - 1) * rowX + entry
    end

    local function lastRowX(entry: number): number
        return (rowZ - 1) * rowX + entry
    end

    local methods = {
        GetNext = getNextEntry,
        GetPrev = getPrevEntry,
        GetForward = getForwardEntry,
        GetBackward = getBackwardEntry
    }

    
    local function fromEntry(index: number): GridMethods
        entryIndex = index

        return methods
    end

    return {
        Nodes = grid,
        Index = fromEntry,
        LastRow = lastRowX,
        FirstRow = firstRowX
    }
end

local function generateRandomPath<grid>(grid: grid & Grid<any>, startIndex: number, endIndex: number): { number }
    local currentIndex = startIndex
    local path = {currentIndex}
    local visited = {}

    visited[currentIndex] = true

    while currentIndex ~= endIndex do
        local directions = {}
        local nextNode = grid.Index(currentIndex):GetNext()
        local prevNode = grid.Index(currentIndex):GetPrev()
        local forwardNode = grid.Index(currentIndex):GetForward()

        if nextNode and not visited[nextNode] then
            table.insert(directions, {node = nextNode, weight = 1.4})
        end
        if prevNode and not visited[prevNode] then
            table.insert(directions, {node = prevNode, weight = 1.4})
        end
        if forwardNode and not visited[forwardNode] then
            table.insert(directions, {node = forwardNode, weight = 1.2})
        end

        for i, direction in directions do
            local distance = math.abs(direction.node - endIndex)
            direction.weight = direction.weight / (distance + 1)
        end


        if #directions > 0 then
            local totalWeight = 0
            for _, direction in directions do
                totalWeight = totalWeight + direction.weight
            end

            local rand = random:NextNumber(0, totalWeight)
            local accumulatedWeight = 0
            local selectedNode

            for _, direction in directions do
                accumulatedWeight = accumulatedWeight + direction.weight
                if rand <= accumulatedWeight then
                    selectedNode = direction.node
                    break
                end
            end

            currentIndex = selectedNode
            table.insert(path, currentIndex)
            visited[currentIndex] = true
        else
            break
        end
    end

    return path
end


return {
    new = new,
    Path = generateRandomPath
}