#include <unistd.h>

#include "BDDTest.h"
#include "Buffer.h"
#include "PubSubClient.h"
#include "ShimClient.h"
#include "trace.h"

byte server[] = {172, 16, 0, 2};

// function declarations
void callback(char* topic, uint8_t* payload, size_t length);
int test_keepalive_pings_idle();
int test_keepalive_pings_with_outbound_qos0();
int test_keepalive_pings_with_inbound_qos0();
int test_keepalive_no_pings_inbound_qos1();
int test_keepalive_disconnects_hung();

void callback(_UNUSED_ char* topic, _UNUSED_ uint8_t* payload, _UNUSED_ size_t length) {
    // handle message arrived
}

int test_keepalive_pings_idle() {
    IT("keeps an idle connection alive (takes 1 minute)");

    ShimClient shimClient;
    shimClient.setAllowConnect(true);

    byte connack[] = {0x20, 0x02, 0x00, 0x00};
    shimClient.respond(connack, 4);

    PubSubClient client(server, 1883, callback, shimClient);
    bool rc = client.connect("client_test1");
    IS_TRUE(rc);

    byte pingreq[] = {0xC0, 0x0};
    shimClient.expect(pingreq, 2);
    byte pingresp[] = {0xD0, 0x0};
    shimClient.respond(pingresp, 2);

    for (int i = 0; i < 50; i++) {
        sleep(1);
        if (i == 15 || i == 31 || i == 47) {
            shimClient.expect(pingreq, 2);
            shimClient.respond(pingresp, 2);
        }
        rc = client.loop();
        IS_TRUE(rc);
    }

    IS_FALSE(shimClient.error());

    END_IT
}

int test_keepalive_pings_with_outbound_qos0() {
    IT("keeps a connection alive that only sends qos0 (takes 1 minute)");

    ShimClient shimClient;
    shimClient.setAllowConnect(true);

    byte connack[] = {0x20, 0x02, 0x00, 0x00};
    shimClient.respond(connack, 4);

    PubSubClient client(server, 1883, callback, shimClient);
    bool rc = client.connect("client_test1");
    IS_TRUE(rc);

    byte publish[] = {0x30, 0xe, 0x0, 0x5, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64};

    for (int i = 0; i < 50; i++) {
        TRACE(i << ":");
        shimClient.expect(publish, 16);
        rc = client.publish("topic", "payload");
        IS_TRUE(rc);
        IS_FALSE(shimClient.error());
        sleep(1);
        if (i == 15 || i == 31 || i == 47) {
            byte pingreq[] = {0xC0, 0x0};
            shimClient.expect(pingreq, 2);
            byte pingresp[] = {0xD0, 0x0};
            shimClient.respond(pingresp, 2);
        }
        rc = client.loop();
        IS_TRUE(rc);
        IS_FALSE(shimClient.error());
    }

    END_IT
}

int test_keepalive_pings_with_inbound_qos0() {
    IT("keeps a connection alive that only receives qos0 (takes 1 minute)");

    ShimClient shimClient;
    shimClient.setAllowConnect(true);

    byte connack[] = {0x20, 0x02, 0x00, 0x00};
    shimClient.respond(connack, 4);

    PubSubClient client(server, 1883, callback, shimClient);
    bool rc = client.connect("client_test1");
    IS_TRUE(rc);

    byte publish[] = {0x30, 0xe, 0x0, 0x5, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64};

    for (int i = 0; i < 50; i++) {
        TRACE(i << ":");
        sleep(1);
        if (i == 15 || i == 31 || i == 47) {
            byte pingreq[] = {0xC0, 0x0};
            shimClient.expect(pingreq, 2);
            byte pingresp[] = {0xD0, 0x0};
            shimClient.respond(pingresp, 2);
        }
        shimClient.respond(publish, 16);
        rc = client.loop();
        IS_TRUE(rc);
        IS_FALSE(shimClient.error());
    }

    END_IT
}

int test_keepalive_no_pings_inbound_qos1() {
    IT("does not send pings for connections with inbound qos1 (takes 1 minute)");

    ShimClient shimClient;
    shimClient.setAllowConnect(true);

    byte connack[] = {0x20, 0x02, 0x00, 0x00};
    shimClient.respond(connack, 4);

    PubSubClient client(server, 1883, callback, shimClient);
    bool rc = client.connect("client_test1");
    IS_TRUE(rc);

    byte publish[] = {0x32, 0x10, 0x0, 0x5, 0x74, 0x6f, 0x70, 0x69, 0x63, 0x12, 0x34, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61, 0x64};
    byte puback[] = {0x40, 0x2, 0x12, 0x34};

    for (int i = 0; i < 50; i++) {
        shimClient.respond(publish, 18);
        shimClient.expect(puback, 4);
        sleep(1);
        rc = client.loop();
        IS_TRUE(rc);
        IS_FALSE(shimClient.error());
    }

    END_IT
}

int test_keepalive_disconnects_hung() {
    IT("disconnects a hung connection (takes 30 seconds)");

    ShimClient shimClient;
    shimClient.setAllowConnect(true);

    byte connack[] = {0x20, 0x02, 0x00, 0x00};
    shimClient.respond(connack, 4);

    PubSubClient client(server, 1883, callback, shimClient);
    bool rc = client.connect("client_test1");
    IS_TRUE(rc);

    byte pingreq[] = {0xC0, 0x0};
    shimClient.expect(pingreq, 2);

    for (int i = 0; i < 32; i++) {
        sleep(1);
        rc = client.loop();
    }
    IS_FALSE(rc);

    int state = client.state();
    IS_TRUE(state == MQTT_CONNECTION_TIMEOUT);

    IS_FALSE(shimClient.error());

    END_IT
}

int main() {
    SUITE("Keep-alive");
    test_keepalive_pings_idle();
    test_keepalive_pings_with_outbound_qos0();
    test_keepalive_pings_with_inbound_qos0();
    test_keepalive_no_pings_inbound_qos1();
    test_keepalive_disconnects_hung();

    FINISH
}
