--[[
A central server to coordinate a rail network. This central server keeps a
graph of the entire network, and handles requests to find paths to route
traffic through.
]]--

local RECEIVE_CHANNEL = 45453

-- ONLY FOR DEBUGGING
-- inspect = require("inspect")
local modem = peripheral.wrap("top") or error("Missing top modem")
modem.open(RECEIVE_CHANNEL)

local function generateStandardNode(id, edgeIds)
    local node = {id = id, connections = {}, type = "JUNCTION"}
    for _, edgeId in pairs(edgeIds) do
        for _, edgeId2 in pairs(edgeIds) do
            if edgeId2 ~= edgeId then
                table.insert(node.connections, {from = edgeId, to = edgeId2})
            end
        end
    end
    return node
end

local function generateStationNode(id, edgeId)
    return {
        id = id,
        connections = {
            {from = nil, to = edgeId},
            {from = edgeId, to = nil}
        },
        type = "STATION"
    }
end

local function loadGraph()
    -- local g = nil
    -- local f = io.open("network_graph.tbl", "r")
    -- g = textutils.unserialize(f:read("*a"))
    -- f:close()
    --return g
    return {
        nodes = {
            generateStandardNode("Junction-HandieVale", {"handievale", "N1", "W1"}),
            generateStandardNode("Junction-Middlecross", {"W1", "N2", "W2", "S1"}),
            generateStandardNode("Junction-Foundry", {"E1", "N3", "W3", "N2"}),
            generateStandardNode("Junction-End", {"E1", "E2", "end"}),
            generateStandardNode("Junction-Klausville", {"N3", "N4", "klausville"}),
            generateStandardNode("Junction-Foundry-West", {"W3", "foundry", "W4"}),
            generateStationNode("station-klausville", "klausville"),
            generateStationNode("station-handievale", "handievale"),
            generateStationNode("station-end", "end"),
            generateStationNode("station-foundry", "foundry")
        },
        edges = {
            {id = "handievale", length = 16},
            {id = "end", length = 48},
            {id = "foundry", length = 45},
            {id = "klausville", length = 12},
            {id = "N1", length = nil},
            {id = "W1", length = 300},
            {id = "N2", length = 600},
            {id = "E1", length = 75},
            {id = "W2", length = nil},
            {id = "S1", length = nil},
            {id = "W3", length = 50},
            {id = "W4", length = nil},
            {id = "N3", length = 350},
            {id = "N4", length = nil}
        }
    }
end

local function filterTable(arr, func)
    local new_index = 1
    local size_orig = #arr
    for old_index, v in ipairs(arr) do
        if func(v, old_index) then
            arr[new_index] = v
            new_index = new_index + 1
        end
    end
    for i = new_index, size_orig do arr[i] = nil end
end

local function findNodeById(graph, nodeId)
    for _, node in pairs(graph.nodes) do
        if node.id == nodeId then return node end
    end
    return nil
end

local function findEdgeById(graph, edgeId)
    for _, edge in pairs(graph.edges) do
        if edge.id == edgeId then return edge end
    end
    return nil
end

local function findEdgeBetweenNodes(graph, fromNode, toNode)
    local edgeIdsFrom = {}
    for _, conn in pairs(fromNode.connections) do
        if conn.to and not tableContains(edgeIdsFrom, conn.to) then
            table.insert(edgeIdsFrom, conn.to)
        end
    end
    local edgeIds = {}
    for _, conn in pairs(toNode.connections) do
        if conn.from and tableContains(edgeIdsFrom, conn.from) and not tableContains(edgeIds, conn.from) then
            table.insert(edgeIds, conn.from)
        end
    end

end

local function tableContains(table, value)
    for _, item in pairs(table) do
        if item == value then return true end
    end
    return false
end

local function removeElement(table, value)
    local idx = nil
    for i, item in pairs(table) do
        if item == value then idx = i break end
    end
    if idx then table.remove(table, idx) end
end

local function findNextEdges(graph, edgeId)
    local edges = {}
    for _, node in pairs(graph.nodes) do
        for _, connection in pairs(node.connections) do
            if connection.from == edgeId then
                table.insert(edges, findEdgeById(connection.to))
            end
        end
    end
    return edges
end

-- Find the set of nodes directly connected to this one via some edges
local function findConnectedNodes(graph, startNode)
    local edges = {}
    local edgeIds = {}
    for _, conn in pairs(startNode.connections) do
        if conn.to ~= nil then
            local edge = findEdgeById(graph, conn.to)
            if edge ~= nil and edge.length ~= nil and not tableContains(edgeIds, edge.id) then
                table.insert(edges, edge)
                table.insert(edgeIds, edge.id)
            end
        end
    end
    local connections = {}
    for _, edge in pairs(edges) do
        for _, node in pairs(graph.nodes) do
            if node.id ~= startNode.id then
                for _, conn in pairs(node.connections) do
                    if conn.from == edge.id then
                        table.insert(connections, {node = node, distance = edge.length, via = edge.id})
                        break
                    end
                end
            end
        end
    end
    return connections
end

local function findPath(graph, startNode, endNode)
    local INFINITY = 1000000000
    local dist = {}
    local prev = {}
    local queue = {}
    for _, node in pairs(graph.nodes) do
        dist[node.id] = INFINITY
        prev[node.id] = nil
        table.insert(queue, node)
    end
    dist[startNode.id] = 0

    while #queue > 0 do
        local minIdx = nil
        local minDist = INFINITY + 1
        for i, node in pairs(queue) do
            if dist[node.id] < minDist then
                minIdx = i
                minDist = dist[node.id]
            end
        end
        if minIdx == nil then return nil end
        local u = table.remove(queue, minIdx)
        if u.id == endNode.id and (prev[u.id] or u.id == startNode.id) then
            local s = {}
            while u ~= nil do
                local via = nil
                local distance = nil
                local node = u
                u = nil
                if prev[node.id] then
                    via = prev[node.id].via
                    distance = prev[node.id].distance
                    u = prev[node.id].node
                end
                table.insert(s, 1, {node = node, via = via, distance = distance})
            end
            return s
        end
        for _, neighbor in pairs(findConnectedNodes(graph, u)) do
            local unvisited = false
            for _, node in pairs(queue) do
                if node.id == neighbor.node.id then
                    unvisited = true
                    break
                end
            end
            if unvisited then
                local alt = dist[u.id] + neighbor.distance
                if alt < dist[neighbor.node.id] then
                    dist[neighbor.node.id] = alt
                    prev[neighbor.node.id] = {node = u, via = neighbor.via, distance = neighbor.distance}
                end
            end
        end
    end
    return nil
end

local function handleRouteRequest(graph, replyChannel, msg)
    if not msg.startNode or not msg.endNode then
        modem.transmit(replyChannel, RECEIVE_CHANNEL, {success = false, error = "Invalid request"})
        return
    end
    local startNode = findNodeById(graph, msg.startNode)
    local endNode = findNodeById(graph, msg.endNode)
    if not startNode or not endNode then
        modem.transmit(replyChannel, RECEIVE_CHANNEL, {success = false, error = "Unknown node(s)"})
        return
    end
    local path = findPath(graph, startNode, endNode)
    if not path then
        modem.transmit(replyChannel, RECEIVE_CHANNEL, {success = false, error = "No valid path"})
        return
    end
    modem.transmit(replyChannel, RECEIVE_CHANNEL, {success = true, route = path})
end

local function handleRequests(graph)
    while true do
        local event, side, channel, replyChannel, msg, dist = os.pullEvent("modem_message")
        if channel == RECEIVE_CHANNEL and msg and msg.command and type(msg.command) == "string" then
            if msg.command == "ROUTE" then
                handleRouteRequest(graph, replyChannel, msg)
            else
                modem.transmit(replyChannel, RECEIVE_CHANNEL, {success = false, error = "Invalid command"})
            end
        end
    end
end

handleRequests(loadGraph())

-- local graph = loadGraph()
-- print("GRAPH:")
-- print(inspect(graph))
-- local startNode = findNodeById(graph, "station-handievale")
-- local endNode = findNodeById(graph, "station-foundry")
-- print("\n\nPATH:")
-- local path = findPath(graph, startNode, endNode)
-- if path then
--     print("Found path!")
--     for i, element in pairs(path) do
--         print(i..". "..element.node.id.." via edge "..inspect(element.via).." @ "..inspect(element.distance))
--     end
-- end