#!/usr/bin/env qore
# -*- mode: qore; indent-tabs-mode: nil -*-

%new-style
%enable-all-warnings
%require-types
%strict-args

%requires ../../../../qlib/Util.qm
%requires ../../../../qlib/QUnit.qm
%requires ../../../../qlib/Logger.qm
%requires ../../../../qlib/HttpServerUtil.qm
%requires ../../../../qlib/HttpServer.qm
%requires ../../../../qlib/ConnectionProvider
%requires ../../../../qlib/WebSocketUtil.qm
%requires ../../../../qlib/WebSocketHandler.qm
%requires ../../../../qlib/WebSocketClient.qm

%exec-class WebSocketHandlerTest

class MyBadWsConnection inherits WebSocketConnection {
    constructor(WebSocketHandler handler) : WebSocketConnection(handler) {
    }

    gotMessage(string msg) {
        send("RESPONSE:"+msg);
    }

    registered() {
        throw "ERROR", "error";
    }
}

class MyWsConnection inherits WebSocketConnection {
    constructor(WebSocketHandler handler) : WebSocketConnection(handler) {
    }

    int gotRawMessage(hash<WsMsgInfo> msg) {
        #printf("raw msg: %y\n", msg);
        return WSC_Process;
    }

    gotMessage(string msg) {
        #printf("msg: %y\n", msg);
        switch (msg) {
            case "STOP":
                stop();
                break;
            case "FLUSH":
                send("flushing");
                flush();
                break;
            case "FRAGMENT":
                send("PART1-", WSOP_Text, False);
                send("PART2", WSOP_Continuation, True);
                break;
            case "PING":
                send("test ping", WSOP_Ping, True);
                break;
            default:
                send("RESPONSE:" + msg);
                break;
        }
    }
}

class MyWsHandler inherits WebSocketHandler {
    public {
        Queue errs();
        bool debug = False;
    }

    constructor(*hash<auto> opts) : WebSocketHandler(NOTHING, opts) {
    }

    WebSocketConnection getConnectionImpl(hash<auto> cx, hash<auto> hdr, string cid) {
        return hdr.bad ? new MyBadWsConnection(self) : new MyWsConnection(self);
    }

    logError(string fmt) {
        if (debug) {
            vprintf("SERVER ERROR " + fmt + "\n", argv);
        }
        errs.push(vsprintf(fmt, argv));
    }

    logInfo(string fmt) {
        if (debug) {
            vprintf("SERVER INFO " + fmt + "\n", argv);
        }
    }

    logDebug(string fmt) {
        if (debug) {
            vprintf("SERVER DEBUG " + fmt + "\n", argv);
        }
    }
}

class MyWsClient inherits WebSocketClient {
    private {
        Counter cnt;
    }

    constructor(code cb, hash<auto> opts, Counter c) : WebSocketClient(cb, opts) {
        cnt = c;
    }

    private handlePing(hash<auto> h) {
        WebSocketClient::handlePing(h);
        # in case of a race condition; only decrement once
        if (cnt.getCount()) {
            cnt.dec();
        }
    }
}

class WebSocketHandlerTest inherits QUnit::Test {
    public {
        HttpServer mServer;
        MyWsHandler wsHandler({"heartbeat": 1s, "heartbeat_msg": "test"});
        int port;
    }

    constructor() : Test("WebSocketHandlerTest", "1.0") {
        addTestCase("WebSocketHandler tests", \webSocketHandlerTests());

        Logger logger("test", LoggerLevel::getLevelInfo());
        if (m_options.verbose > 2) {
            logger.addAppender(new StdoutAppender());
        }
        hash<HttpServerOptionInfo> http_opts = <HttpServerOptionInfo>{
            "logger": logger,
            "debug": True,
        };
        mServer = new HttpServer(http_opts);
        mServer.setHandler("ws-handler", "", NOTHING, wsHandler);
        mServer.setDefaultHandler("ws-handler", wsHandler);
        port = mServer.addListener(<HttpListenerOptionInfo>{"service": 0}).port;

        set_return_value(main());

        mServer.stop();
    }

    webSocketHandlerTests() {
        assertThrows("WEBSOCKETHANDLER-OPTION-ERROR", sub () { new MyWsHandler({"heartbeat": {}}); });
        assertThrows("WEBSOCKETHANDLER-OPTION-ERROR", sub () { new MyWsHandler({"heartbeat": -1}); });

        {
            MyWsHandler wsh();
            assertEq(20, wsh.getHeartbeat());
            assertEq("heartbeat", wsh.getHeartbeatMessage());
        }

        assertEq(1, wsHandler.getHeartbeat());
        assertEq("test", wsHandler.getHeartbeatMessage());

        {
            Counter c(1);
            MyWsClient wsc(sub (*data msg) {}, {"url": "ws://localhost:" + port}, c);
            wsc.connect();
            c.waitForZero();
            assertTrue(wsc.isOpen());
        }

        # check bad handler
        {
            WebSocketClient wsc(sub (*data msg) {}, {"url": "ws://localhost:" + port});
            wsc.connect({"hdr": {"Bad": 1}});
            wsc.disconnect();
        }

        # check close connection from server side
        {
            Counter c(1);
            wsHandler.setHeartbeat(20);  # to not decrement counter
            MyWsClient wsc(sub (*data msg) {c.dec(); }, {"url": "ws://localhost:" + port}, c);
            wsc.connect();
            assertTrue(wsc.isOpen());
            wsc.send("STOP");
            c.waitForZero();
        }

        # check fragmented messaging and flushing I/O
        {
            Counter c(1);
            Queue msgs();
            code callback = sub (*data msg) {
                msgs.push(msg);
            };
            list<hash<WsMsgInfo>> raw_msgs;
            code raw_callback = int sub (hash<WsMsgInfo> msg) {
                raw_msgs += msg;
                return WSC_Process;
            };
            code log = code sub (string str) {
                return sub (string fmt) {
                    #vprintf("CLIENT " + str + ": " + fmt + "\n", argv);
                    remove argv;
                };
            };
            MyWsClient wsc(callback, {
                "url": "ws://localhost:" + port,
                "raw_callback": raw_callback,
                "log": log("INFO"),
                "errlog": log("ERROR"),
                "debuglog": log("DEBUG"),
            }, c);
            wsc.connect();
            assertTrue(wsc.isOpen());
            wsc.send("FRAG", WSOP_Text, False);
            wsc.send("MENT", WSOP_Continuation, True);
            assertEq("PART1-PART2", msgs.get(5s));
            wsc.send("FLUSH");
            assertEq("flushing", msgs.get(5s));
            wsc.send("PING");
            # wait for ping
            c.waitForZero();
            assertEq(WSOP_Text, raw_msgs[0].op);
            assertEq(False, raw_msgs[0].fin);
            assertEq(WSOP_Continuation, raw_msgs[1].op);
            assertEq(True, raw_msgs[1].fin);
        }

        wsHandler.errs.get(10s);
        assertEq(0, wsHandler.errs.size());
    }

    private log(string str) {
        if (m_options.verbose > 2)
            vprintf(str + "\n", argv);
    }
}
