#include "common.h"


extern "C" {
#include "prrt/proto/types/block.h"
#include "prrt/proto/types/codingParams.h"
#include "prrt/proto/types/packet.h"
#include "prrt/util/dbg.h"
#include "prrt/proto/vdmcode/block_code.h"
#include "prrt/util/common.h"
}

class PrrtBlockTest : public ::testing::Test {
protected:
    virtual void SetUp()
    {
    }

    virtual void TearDown() {
    }

};

TEST_F(PrrtBlockTest, VDMCode)
{
    int i, j, q;
    uint8_t k = 4;
    uint8_t n = 7;
    uint8_t r = n - k;
    uint8_t length = 3;


    PrrtCodingConfiguration* params = PrrtCodingConfiguration_create(k, n, 1, &r);
    PrrtCoder *coder = PrrtCoder_create(params);

    for(q = 0; q < 4; q++) {
        gf **src = (gf **) calloc(k, sizeof(gf *));
        src[0] = (gf *) calloc(length, sizeof(gf));
        src[0][0] = 1; src[0][2] = 3;
        src[1] = (gf *) calloc(length, sizeof(gf));
        src[1][0] = 2; src[1][2] = 7;
        src[2] = (gf *) calloc(length, sizeof(gf));
        src[2][0] = 3; src[2][2] = 6;
        src[3] = (gf *) calloc(length, sizeof(gf));
        src[3][0] = 4; src[3][1]= 9; src[3][2] = 5;

        gf **encFec = (gf **) calloc(1, sizeof(gf *) * r);
        for(j = 0; j < r; j++) {
            encFec[j] = (gf *) calloc(length, sizeof(gf));
            PrrtCoder_encode(coder, src, encFec[j], j + k, length);
        }

        gf **decFec = (gf **) calloc(1, sizeof(gf *) * k);
        int16_t *idx_p = (int16_t *) calloc(k, sizeof(int16_t));

        for(i = 0; i < k; i++) {
            decFec[i] = (gf *) calloc(length, sizeof(gf));
        }

        memcpy(decFec[0], src[0], length);
        idx_p[0] = 0;

        memcpy(decFec[1], src[1], length);
        idx_p[1] = 1;

        memcpy(decFec[2], encFec[0], length);
        idx_p[2] = 4;

        memcpy(decFec[3], src[3], length);
        idx_p[3] = 3;

        PrrtCoder_decode(coder, decFec, idx_p, length);

        ASSERT_EQ(1, decFec[0][0]);
        ASSERT_EQ(0, decFec[0][1]);
        ASSERT_EQ(3, decFec[0][2]);

        ASSERT_EQ(2, decFec[1][0]);
        ASSERT_EQ(0, decFec[1][1]);
        ASSERT_EQ(7, decFec[1][2]);

        ASSERT_EQ(3, decFec[2][0]);
        ASSERT_EQ(0, decFec[2][1]);
        ASSERT_EQ(6, decFec[2][2]);

        ASSERT_EQ(4, decFec[3][0]);
        ASSERT_EQ(9, decFec[3][1]);
        ASSERT_EQ(5, decFec[3][2]);

        // free allocated memory

        for (i = 0; i < k; ++i) {
            free(decFec[i]);
        }
        free(idx_p);
        free(decFec);

        for (j = 0; j < r; ++j) {
            free(encFec[j]);
        }
        free(encFec);

        free(src[3]);
        free(src[2]);
        free(src[1]);
        free(src[0]);
        free(src);
    }

    PrrtCoder_destroy(coder);
    PrrtCodingConfiguration_destroy(params);
}

TEST_F(PrrtBlockTest, EncodeDecode)
{
    uint8_t r = 3;
    PrrtCodingConfiguration *cpar = PrrtCodingConfiguration_create(4, 7, 1, &r);
    PrrtCoder *coder = PrrtCoder_create(cpar);

    PrrtBlock *encBlock = PrrtBlock_create(PrrtCodingConfiguration_copy(cpar), PrrtCoder_copy(coder), 1);
    PrrtBlock *decBlock = PrrtBlock_create(PrrtCodingConfiguration_copy(cpar), PrrtCoder_copy(coder), 1);

    prrtSequenceNumber_t base = 1;

    PrrtPacket *packets[4];
    PrrtPacket *refPackets[4];

    PrrtPacket *redPackets[4];

    for(uint32_t i = 0; i < 4; i++) {
        char data[3];
        sprintf(data, "%d", i);
        if(i == 3) {
            // Last packet has a different size.
            sprintf(data, "D:%d", i);
        }
        packets[i] = PrrtPacket_create_data_packet(0, data, (prrtPacketLength_t) strlen(data), i + 1, 0);
        packets[i]->index = (uint8_t) i;
        refPackets[i] = PrrtPacket_create_data_packet(0, data, (prrtPacketLength_t) strlen(data), i + 1, 0);
        ASSERT_TRUE(PrrtBlock_insert_data_packet(encBlock, packets[i]));
    }

    // ENCODING
    ASSERT_TRUE(PrrtBlock_encode_ready(encBlock));
    uint32_t j = 0;
    PrrtBlock_encode(encBlock, &base);

    uint32_t pkt_count = List_count(encBlock->dataPackets);
    for(j = 0; j < pkt_count; j++) {
        PrrtPacket *data_pkt = PrrtBlock_get_first_data(encBlock);
        ASSERT_EQ(refPackets[j]->sequenceNumber, data_pkt->sequenceNumber);
    }
    uint32_t red_count = List_count(encBlock->redundancyPackets);
    for(j = 0; j < red_count; j++) {
        PrrtPacket *red_pkt = PrrtBlock_get_first_red_data(encBlock);
        redPackets[j] = red_pkt;
    }

    printf("------------------------\n");

    // DECODING
    PrrtBlock_insert_redundancy_packet(decBlock, redPackets[0]);
    PrrtBlock_insert_redundancy_packet(decBlock, redPackets[1]);
    PrrtBlock_insert_redundancy_packet(decBlock, redPackets[2]);
    PrrtBlock_insert_data_packet(decBlock, PrrtPacket_copy(packets[3]));

    ASSERT_TRUE(PrrtBlock_decode_ready(decBlock));

    PrrtBlock_decode(decBlock);

    for(int k = 0; k < 4; ++k) {
        PrrtPacket *ptr = PrrtBlock_get_first_data(decBlock);
        prrtIndex_t idx = static_cast<prrtIndex_t>(ptr->sequenceNumber - 1);

        const char *s1 = (const char *) ((char *) refPackets[idx]->payload + PRRT_PACKET_DATA_HEADER_SIZE);
        const char *s2 = (const char *) ((char *) ptr->payload + PRRT_PACKET_DATA_HEADER_SIZE);

        prrtPacketLength_t refLength = ((PrrtPacketDataPayload*) refPackets[idx]->payload)->dataLength;
        prrtPacketLength_t ptrLength = ((PrrtPacketDataPayload*) ptr->payload)->dataLength;

        ASSERT_STREQ(s1, s2);
        ASSERT_EQ(refLength, ptrLength);

        PrrtPacket_destroy(ptr);
    }

    PrrtBlock_destroy(encBlock);
    PrrtBlock_destroy(decBlock);
    PrrtCoder_destroy(coder);
    PrrtCodingConfiguration_destroy(cpar);

    for (auto packet: packets) {
        PrrtPacket_destroy(packet);
    }

    for (auto packet: refPackets) {
        PrrtPacket_destroy(packet);
    }
}

TEST_F(PrrtBlockTest, NoCodingConfigured)
{
    PrrtCodingConfiguration *cpar = PrrtCodingConfiguration_create(1, 1, 0, NULL);
    PrrtCoder* coder = PrrtCoder_create(cpar);
    PrrtBlock *encBlock = PrrtBlock_create(PrrtCodingConfiguration_copy(cpar), PrrtCoder_copy(coder), 1);

    prrtSequenceNumber_t base = 1;

    PrrtPacket *packets[1];
    PrrtPacket *refPackets[1];

    for(uint32_t i = 0; i < 1; i++) {
        char data[3];
        sprintf(data, "%d", i);
        packets[i] = PrrtPacket_create_data_packet(0, data, (prrtPacketLength_t) strlen(data) + 1, i + 1, 0);
        packets[i]->index = (uint8_t) i;
        refPackets[i] = PrrtPacket_create_data_packet(0, data, (prrtPacketLength_t) strlen(data) + 1, i + 1, 0);
        ASSERT_TRUE(PrrtBlock_insert_data_packet(encBlock, packets[i]));
    }

    // ENCODING
    ASSERT_TRUE(PrrtBlock_encode_ready(encBlock));
    uint32_t j = 0;
    PrrtBlock_encode(encBlock, &base);

    uint32_t pkt_count = List_count(encBlock->dataPackets);
    for(j = 0; j < pkt_count; j++) {
        PrrtPacket *data_pkt = PrrtBlock_get_first_data(encBlock);
        ASSERT_EQ(refPackets[j]->sequenceNumber, data_pkt->sequenceNumber);
    }
    uint32_t red_count = List_count(encBlock->redundancyPackets);
    ASSERT_EQ(red_count, 0);

    for (int i = 0; i < 1; ++i) {
        PrrtPacket_destroy(packets[i]);
        PrrtPacket_destroy(refPackets[i]);
    }

    PrrtBlock_destroy(encBlock);
    PrrtCoder_destroy(coder);
    PrrtCodingConfiguration_destroy(cpar);
}