TURTI_WebsocketProtocol

Protocol = {}
Protocol.__index = Protocol

ProtocolWebsocket = {}
ProtocolWebsocket.__index = ProtocolWebsocket

function Protocol:new()
    local protocol = {}
    setmetatable(protocol, self)
    protocol.__index = protocol
    return protocol
end

function Protocol:setOnOpen(f)
    self.onOpen = f
end

function Protocol:setOnClose(f)
    self.onClose = f
end

function Protocol:setOnMessage(f)
    self.onMessage = f
end

function Protocol:send(msg)
    self.exSend(msg)
end

function ProtocolWebsocket:new(url, protocol)
    local protocolWebsocket = {}
    setmetatable(protocolWebsocket, self)
    protocolWebsocket.__index = protocolWebsocket
    protocolWebsocket.url = url .. "?computer_id=" .. os.getComputerID()
    protocolWebsocket.startProtocol = protocol
    protocolWebsocket.protocolChangeHandshakeRunning = false
    protocolWebsocket.protocolChangeHandshakeDoneReceived = true
    protocolWebsocket.protocolChangeHandshakeClientDone = false
    return protocolWebsocket
end

function ProtocolWebsocket:close()
    self.alive = false
    self.websocket.close()
end

function ProtocolWebsocket:start(unsupportedEventHandler)
    local websocket, err = http.websocket(self.url)
    if not websocket then
        error(err)
    end
    self.websocket = websocket;
    self.alive = true
    self:setProtocol(self.startProtocol)
    if self.protocol.onOpen then
        self.protocol.onOpen()
    end
    while self.alive do
        local type, url, msg, arg = os.pullEvent()
        if type == "websocket_message" then
            if url == self.url then
                self:handleWebsocketMessage(msg)
            end
        elseif type == "websocket_closed" then
            if self.protocol.onClose then
                self.alive = false
                self.protocol.onClose(tonumber(arg), msg)
            end
        else
            if unsupportedEventHandler then
                unsupportedEventHandler(type, url, msg, arg)
            end
        end
    end
end

function ProtocolWebsocket:handleWebsocketMessage(msg)
    if msg:sub(1, 1) == '&' then
        if msg:sub(2, 2) == '0' then
            if self.protocolChangeHandshakeRunning then
                error("io isn't available during protocol change handshake");
            elseif not self.protocolChangeHandshakeDoneReceived then
                self.protocolChangeHandshakeDoneReceived = true;
                if self.protocol.onOpen then
                    self.protocol.onOpen()
                end
            end
            if self.protocol.onMessage then
                self.protocol.onMessage(msg:sub(3, #msg))
            end
        elseif msg:sub(2, 2) == '1' then
            msg = msg:sub(3, #msg)
            if msg:sub(1, 15) == "switchProtocol:" then
                local type = tonumber(msg.sub(16, #msg));
                self.protocolChangeHandshakeRunning = true;
                self.protocolChangeHandshakeDoneReceived = false;
                if self.protocol.onClose then
                    self.alive = false
                    self.protocol.onClose(1000, "protocol change")
                end
                if not self.switchProtocol then
                    error("protocol-switch handler 'onSwitchProtocol' is not set")
                end
                self.setProtocol(self.onSwitchProtocol(type));
                self.sendInternalMessage("getProtocolDone");
                self.protocolChangeHandshakeRunning = false;
            end
        elseif msg == "getProtocolDone" then
            self.protocolChangeHandshakeClientDone = true;
        elseif msg == "switchProtocolDone" then
            if not self.protocolChangeHandshakeDoneReceived then
                self.protocolChangeHandshakeDoneReceived = true;
                if self.protocol.onOpen then
                    self.protocol.onOpen()
                end
            end
        else
            error("invalid message" .. msg)
        end
    else
        error("invalid message" .. msg)
    end
end

function ProtocolWebsocket:sendInternalMessage(msg)
    self.websocket:send("&1" + msg);
end

function ProtocolWebsocket:setProtocol(protocol)
    if self.protocol then
        self.protocol.exSend = nil
    end
    self.protocol = protocol
    protocol.exSend = function(data)
        if (self.protocolChangeHandshakeRunning) then
            error("io isn't available during protocol change handshake")
        end
        self.websocket.send("&0" .. data)
    end
end

ActionProtocol = Protocol:new()

function ActionProtocol:new()
    local ap = Protocol:new()
    setmetatable(ap, self)
    ap.__index = self
    ap.answers = {}
    ap.currentActionId = 0
    ap.argumentStorage = {}
    ap.argumentStorageSize = 0
    ap:setOnMessage(function(msg)
        if msg:sub(1, 1) ~= "@" then
            error("invalid message: " .. msg)
        end
        if msg:sub(2, 2) == "0" then
            local action = splitComma(msg, 3)
            if action[1] == "_newArgument" then
                table.insert(ap.argumentStorage, action[2])
                ap.argumentStorageSize = ap.argumentStorageSize + 1
            elseif action[1] == "_addToLastArgument" then
                ap.argumentStorage[ap.argumentStorageSize] = ap.argumentStorage[ap.argumentStorageSize] .. action[2]
            else
                if ap.onAction then
                    ap.onAction(concatTables(
                            concatTables({ action[1] }, ap.argumentStorage),
                            subArray(action, 2)))
                end
                ap.argumentStorage = {}
                ap.argumentStorageSize = 0
            end
        elseif msg:sub(2, 2) == "1" then
            local action = splitComma(msg, 3)
            local actionId = tonumber(action[1])
            local answer
            if ap.onAction then
                answer = ap.onAction(subArray(action, 2))
            end
            if not answer then
                answer = {}
            end
            ap:send(createMessage("2" .. actionId .. ",", answer))
        elseif msg:sub(2, 2) == "2" then
            local action = splitComma(msg, 3)
            local actionId = tonumber(action[1])
            ap.answers[actionId] = subArray(action, 2)
            else
            print("unknownMessage: "..msg)
        end
    end)
    return ap
end

function concatTables(t1, t2)
    if not t1 then
        return t2
    elseif not t2 then
        return t1
    end
    for _, v in ipairs(t2) do
        table.insert(t1, v)
    end
    return t1
end

function subArray(array, from)
    local result = {}
    for i, v in ipairs(array) do
        if i >= from then
            table.insert(result, v)
        end
    end
    return result
end

function splitComma(msg, start)
    local action = {}
    local current = ""
    local escaped = false
    for i = start, #msg do
        local c = msg:sub(i, i)
        if not escaped then
            if c == "\\" then
                escaped = true
            elseif c == "," then
                table.insert(action, current)
                current = ""
            else
                current = current .. c
            end
        else
            current = current .. c
            escaped = false
        end
    end
    table.insert(action, current)
    return action
end

function ActionProtocol:setOnAction(f)
    self.onAction = f
end

function createMessage(prefix, args)
    local result = "@" .. prefix
    for i, v in ipairs(args) do
        if i ~= 1 then
            result = result .. ","
        end
        result = result .. escapeInMessage(v, ",", "\\")
    end
    return result
end

function ActionProtocol:execute(...)
    self:send(createMessage("0", { ... }))
end

function ActionProtocol:executeGetAnswer(...)
    local id = self.currentActionId
    self.currentActionId = self.currentActionId + 1
    self:send(createMessage("1" .. id .. ",", { ... }))
    while true do
        local answer = self.answers[id]
        if answer then
            return answer
        end
        coroutine.yield()
    end
end

function escapeInMessage(msg, ch, chEscape)
    local result = ""
    for i = 1, #msg do
        local c = msg:sub(i, i)
        if c == ch or c == chEscape then
            result = result .. "\\"
        end
        result = result .. c
    end
    return result
end