Railcraft Server

--Server Program:
--  Tracks trains across network.
--  Handles requests by load balancing.
--  Keeps a list of stations.
--Constants:
rednet.open("top")
--Files to store permanent or sensitive data.
local TRAINFILE = "trains"
local STATIONFILE = "stations"
--Protocols:
local METADATA_PROT = "md"
local METADATA_RQ = "rq"
local UPDATE_PROT = "ud"
local UPDATE_RQ = "rq"
local DISPATCH_PROT = "dp"
local REQUEST_PROT = "rq"
local STATION_LIST_PROT = "sl"
local STATION_LIST_RQ = "rq"
local REBOOT_PROT = "rb"
local REBOOT_RQ = "rb"
--Stations List:
local stations = {}
--Trains List:
local trains = {}

--Returns the number of items in a table.
function tableSize(tbl)
    local count = 0
    for i in pairs(tbl) do count = count + 1 end
    return count
end

--Loads the trainslist into memory.
function loadTrainsFromFile()
    if (not fs.exists(TRAINFILE)) then
        local f = fs.open(TRAINFILE, "w")
        f.write("{}")
        f.close()
    end
    local f = fs.open(TRAINFILE, "r")
    trains = textutils.unserialize(f.readAll())
    f.close()
end

--Saves the trainslist into the file.
function saveTrainsToFile()
    local f = fs.open(TRAINFILE, "w")
    f.write(textutils.serialise(trains))
    f.close()
end

--Loads the stations list from the file.
function loadStationsFromFile()
    if (not fs.exists(STATIONFILE)) then
        local f = fs.open(STATIONFILE, "w")
        f.write("{}")
        f.close()
    end
    local f = fs.open(STATIONFILE, "r")
    stations = textutils.unserialize(f.readAll())
    f.close()
end

--Saves the stations list to the file.
function saveStationsToFile()
    local f = fs.open(STATIONFILE, "w")
    f.write(textutils.serialize(stations))
    f.close()
end

--Gets a station id from the dest.
function getStationId(dest)
    for i in pairs(stations) do
        if (stations[i].dest == dest) then
            return stations[i].id
        end
    end
end

--Returns a table of all trains at, or approaching a station.
function getTrainsAtStation(stationId)
    local trainsAtStation = {}
    for i in pairs(trains) do
        if (trains[i].loc == stations[stationId].dest or trains[i].dest == stations[stationId].dest) then
            table.insert(trainsAtStation, trains[i])
        end
    end
    return trainsAtStation
end

--Returns the list of trains parked at a station.
function getTrainsParkedAtStation(stationId)
    local trainsParked = {}
    local trainsAtStation = getTrainsAtStation(stationId)
    for i in pairs(trainsAtStation) do
        if (trainsAtStation[i].loc == stations[stationId].dest) then
            table.insert(trainsParked, trainsAtStation[i])
        end
    end
    return trainsParked
end

--Broadcasts that all stations respond with metadata.
function sendMetadataRequest()
    term.setTextColor(colors.blue)
    rednet.broadcast(METADATA_RQ, METADATA_PROT)
    print("[M]Sent Metadata request.")
end

--Received update from a station containing all trains at that station.
function onUpdateReceived(msg, stationId)
    if (stations[stationId] == nil) then
        rednet.send(stationId, METADATA_RQ, METADATA_PROT)
        term.setTextColor(colors.blue)
        print("[M]Sent Metadata request.")
        return
    end
    term.setTextColor(colors.lightBlue)
    print("[U]Update received from station: "..stations[stationId].dest)
    for k in pairs(msg) do
        trains[msg[k]] = {id = msg[k], loc = stations[stationId].dest, dest = ""}
        saveTrainsToFile()
    end
end

--Received metadata from a station.
function onMetadataReceived(msg)
    term.setTextColor(colors.blue)
    print("[M]Metadata received from: "..msg.dest)
    stations[msg.id] = msg
    saveStationsToFile()
end

--Prints all trains and their destinations.
function printTrains()
    term.setTextColor(colors.orange)
    print("Trains: ")
    for i in pairs(trains) do
        print("  Train "..i..": Dest: "..trains[i].dest.." Loc: "..trains[i].loc)
    end
end

--Ensure trains are spread out as much as possible.
function balanceTrains()
    term.setTextColor(colors.pink)
    print("[BAL]Balancing start {")
    for i in pairs(stations) do
        local trainsAtStation = getTrainsAtStation(stations[i].id)
        if (tableSize(trainsAtStation) == 0) then
        --No trains at this station, so find an eligible train to send to it.
            term.setTextColor(colors.magenta)
            print("*[BAL]No trains at: "..stations[i].dest)
            for k in pairs(stations) do
                local trainsParked = getTrainsParkedAtStation(stations[k].id)
                if (tableSize(trainsParked) > 1) then
                --This station has a train to spare, so send it to the needy station.
                    sendTrainTo(trainsParked[1].id, stations[i].dest)
                    break
                end
            end
        elseif (tableSize(trainsAtStation) > stations[i].trackCount) then
        --Station is over-saturated, so find a station that can take a train.
            term.setTextColor(colors.magenta)
            print("*[BAL]Station "..stations[i].dest.." is saturated.")
            for k in pairs(stations) do
                local trainsAtPotentialStation = getTrainsAtStation(stations[k].id)
                if (tableSize(trainsAtPotentialStation) < stations[k].trackCount) then
                    --Station is below capacity, so it can take a train.
                    local trainsParked = getTrainsParkedAtStation(stations[i].id)
                    if (tableSize(trainsParked) > 0) then
                    --Check that there is at least one train parked at the saturated station.
                        sendTrainTo(trainsParked[1].id, stations[k].dest)
                    end
                end
            end
        end
    end
    term.setTextColor(colors.pink)
    print("[BAL]Balancing done.")
end

--Attempts to send a train from its current station to a new one.
function sendTrainTo(trainId, dest)
    rednet.send(getStationId(trains[trainId].loc), {dest, trainId}, DISPATCH_PROT)
    trains[trainId].loc = ""
    trains[trainId].dest = dest
    saveTrainsToFile()
    term.setTextColor(colors.lime)
    print("  Sent train "..trainId.." to "..dest)
end

--Initialize the server.
function init()
    term.clear()
    term.setCursorPos(1,1)
    term.setTextColor(colors.yellow)
    print("Initializing server. ID: "..os.getComputerID())
    loadStationsFromFile()
    loadTrainsFromFile()
    sendMetadataRequest()
    broadcastStationList()
end

--Deals with station-initiated requests for train dispatch.
function onRequestReceived(msg, stationId)
    local parkedTrains = getTrainsParkedAtStation(stationId)
    if (tableSize(parkedTrains) < 1) then
        return
    else
        sendTrainTo(parkedTrains[1].id, msg)
        balanceTrains()
    end
end

--Sends the requester a copy of the stations list.
function onStationListRequestReceived(stationId)
    local stationList = {}
    for i in pairs(stations) do
        table.insert(stationList, {dest = stations[i].dest, name = stations[i].name})
    end
    rednet.send(stationId, stationList, STATION_LIST_PROT)
end

--Broadcasts the stations list to all listening stations.
function broadcastStationList()
    local stationList = {}
    for i in pairs(stations) do
        table.insert(stationList, {dest = stations[i].dest, name = stations[i].name})
    end
    rednet.broadcast(stationList, STATION_LIST_PROT)
end

--Forces a reboot of all stations.
function broadcastRebootRequest()
    rednet.broadcast(REBOOT_RQ, REBOOT_PROT)
    term.setTextColor(colors.red)
    print("Sent reboot request.")
    os.sleep(2)
end

--Main event handling function.
function eventHandler()
    local event,p1,p2,p3,p4,p5 = os.pullEvent()
    if (event == "rednet_message") then
        if (p3 == UPDATE_PROT) then
            onUpdateReceived(p2, p1)
        elseif (p3 == METADATA_PROT) then
            onMetadataReceived(p2)
        elseif (p3 == REQUEST_PROT) then
            onRequestReceived(p2, p1)
        elseif (p3 == STATION_LIST_PROT and p2 == STATION_LIST_RQ) then
            onStationListRequestReceived(p1)
        end
    elseif (event == "key") then
        if (p1 == keys.r) then
            --Send reboot.
            broadcastRebootRequest()
        end
    end
end

--Mainloop:
init()
while (true) do
    eventHandler()
end