#include <gtest/gtest.h>

#include <inflate.h>

/* Test that we read the bits in the correct order */ 
TEST (BufferTests, GetNextBitTest) {
    reset_pos();

    /* The 4-byte buffer in memory */ 
    unsigned char buf[4] = {175,  /* = 10101111 */ 
                            240,  /* = 11110000 */ 
                            15,   /* = 00001111 */ 
                            204}; /* = 11001100 */

    /* 
     * The real stream order. 
     * If the bytes are laid out in memory as 
     * [b0b1..b7] [b8...b15] ...
     * Then the correct stream order is:
     * [b7...b0] [b15...b8] ...
     */
    int stream_order[32] = {1, 1, 1, 1, 0, 1, 0, 1,
                            0, 0, 0, 0, 1, 1, 1, 1,
                            1, 1, 1, 1, 0, 0, 0, 0,
                            0, 0, 1, 1, 0, 0, 1, 1};

    for (int i = 0; i < 32; i++) {
        // TODO this cast should probably not be necessary
        // change type signature to unsigned char in inflate
        ASSERT_EQ (get_next_bit((char *) buf), stream_order[i]);
    }
}

/* 
 * Test that we get the correct numerical value when 
 * reading from a buffer. 
 * In this test, all integers are interpreted in 
 * STREAM ORDER. That is, if the bitstream 
 * is (b0, b1, b2, ..., bn), we interpret 
 * b0 as the MSB and bn as the LSB. Note that the order
 * of the bitstream is not necessarily the order that 
 * the bits are laid out in memory; see get_next_bit().
 */
TEST (BufferTests, StreamOrderTest) {
    reset_pos();

    unsigned char buf[2] = {220,  /* = 11011100 */
                            145}; /* = 10010001 */

    /* First 3 bits are 001 */
    ASSERT_EQ (get_n_bits((char *) buf, 3, false), 1);
    /* Next 5 are 11011 */ 
    ASSERT_EQ (get_n_bits((char *) buf, 5, false), 27);
    /* Next 2 are 10 */
    ASSERT_EQ (get_n_bits((char *) buf, 2, false), 2);
    /* Next 6 are 001001 */
    ASSERT_EQ (get_n_bits((char *) buf, 6, false), 9); 
}

/* 
 * Like the test above, but interprets integers in REVERSE STREAM ORDER.
 * That is, if we read in the bitstream (b0, b1, ..., bn), 
 * b0 is the LSB and bn is the MSB. 
 */
TEST (BufferTests, RevStreamOrderTest) {
    reset_pos();

    unsigned char buf[2] = {220,  /* = 11011100 */
                            145}; /* = 10010001 */

    /* First 3 bits are 001 */
    ASSERT_EQ (get_n_bits((char *) buf, 3, true), 4);
    /* Next 5 are 11011 */ 
    ASSERT_EQ (get_n_bits((char *) buf, 5, true), 27);
    /* Next 2 are 10 */
    ASSERT_EQ (get_n_bits((char *) buf, 2, true), 1);
    /* Next 6 are 001001 */
    ASSERT_EQ (get_n_bits((char *) buf, 6, true), 36); 
}

/* Test making a huffman tree from a sequence of lengths */ 
TEST (HuffmanTests, TestMakeHuffman) {
    int n_symbols = 8; 
    int lens[8] = {3, 3, 3, 3, 3, 2, 4, 4}; 

    huffman_t *hf = make_huffman(lens, n_symbols); 
    
    int bl_counts[MAX_LENGTH + 1] = {0, 0, 1, 5, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; 

    /* Check bl_counts */
    for (int i = 0; i < MAX_LENGTH + 1; i++) { 
        ASSERT_EQ (hf->bl_counts[i], bl_counts[i]);     
    }

    /* 
     * Check min_codes for entries with non-zero bl_count 
     * It doesn't really matter what the other entries are. 
     */ 
    ASSERT_EQ (hf->min_codes[2], 0); 
    ASSERT_EQ (hf->min_codes[3], 2);
    ASSERT_EQ (hf->min_codes[4], 14);  

    /* Check that the alphabet arrays were populated correctly. */ 
    ASSERT_EQ (hf->alphabet[2][0], 5);
    for (int i = 0; i < 5; i++) ASSERT_EQ (hf->alphabet[3][i], i); 
    ASSERT_EQ (hf->alphabet[4][0], 6);
    ASSERT_EQ (hf->alphabet[4][1], 7);

    destroy_huffman(hf); 
}

/* 
 * Test that make_huffman does not assign a code to symbols with 
 * a length of 0. 
 */
TEST (HuffmanTests, TestMakeHuffmanZeroLens) { 
    int n_symbols = 8; 
    int lens[8] = {1, 0, 2, 3, 3, 0, 4, 3};

    huffman_t *hf = make_huffman(lens, n_symbols); 
    
    int bl_counts[MAX_LENGTH + 1] = {0, 1, 1, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; 
    
    /* Check bl_counts */
    for (int i = 0; i < MAX_LENGTH + 1; i++) {
        ASSERT_EQ (hf->bl_counts[i], bl_counts[i]); 
    }

    /* Check min_codes */
    ASSERT_EQ (hf->min_codes[1], 0);
    ASSERT_EQ (hf->min_codes[2], 2);
    ASSERT_EQ (hf->min_codes[3], 6);
    ASSERT_EQ (hf->min_codes[4], 18);

    /* Check alphabet. */
    ASSERT_EQ (hf->alphabet[1][0], 0);

    /* 1 and 5 are not used! (length = 0) */
    ASSERT_EQ (hf->alphabet[2][0], 2);

    ASSERT_EQ (hf->alphabet[3][0], 3);
    ASSERT_EQ (hf->alphabet[3][1], 4); 
    ASSERT_EQ (hf->alphabet[3][2], 7); 

    ASSERT_EQ (hf->alphabet[4][0], 6); 

    destroy_huffman(hf); 
}

/* 
 * Test reading a sequence of literals/distances from a buffer. 
 * In this test, the FIXED Huffman codes are used. 
 */
TEST (HuffmanTests, TestReadChunkFixed) {

    /*
     * This is the mapping of literal or length -> code. 
     * The DEFLATE specification defines how to map the literals 
     * (0, 1, ..., 287) to their codes. However, these codes are to 
     * be interpreted in stream order. The numbers below are the actual
     * memory representation of each code, not their stream-order 
     * numerical value.  
     */
    int codes[288] = {12, 140, 76, 204, 44, 172, 108, 236, 28, 156, 92, 220, 
                      60, 188, 124, 252, 2, 130, 66, 194, 34, 162, 98, 226, 18, 
                      146, 82, 210, 50, 178, 114, 242, 10, 138, 74, 202, 42, 
                      170, 106, 234, 26, 154, 90, 218, 58, 186, 122, 250, 6, 
                      134, 70, 198, 38, 166, 102, 230, 22, 150, 86, 214, 54, 
                      182, 118, 246, 14, 142, 78, 206, 46, 174, 110, 238, 30, 
                      158, 94, 222, 62, 190, 126, 254, 1, 129, 65, 193, 33, 
                      161, 97, 225, 17, 145, 81, 209, 49, 177, 113, 241, 9, 
                      137, 73, 201, 41, 169, 105, 233, 25, 153, 89, 217, 57, 
                      185, 121, 249, 5, 133, 69, 197, 37, 165, 101, 229, 21, 
                      149, 85, 213, 53, 181, 117, 245, 13, 141, 77, 205, 45, 
                      173, 109, 237, 29, 157, 93, 221, 61, 189, 125, 253, 19, 
                      275, 147, 403, 83, 339, 211, 467, 51, 307, 179, 435, 115, 
                      371, 243, 499, 11, 267, 139, 395, 75, 331, 203, 459, 43, 
                      299, 171, 427, 107, 363, 235, 491, 27, 283, 155, 411, 91, 
                      347, 219, 475, 59, 315, 187, 443, 123, 379, 251, 507, 7, 
                      263, 135, 391, 71, 327, 199, 455, 39, 295, 167, 423, 103, 
                      359, 231, 487, 23, 279, 151, 407, 87, 343, 215, 471, 55, 
                      311, 183, 439, 119, 375, 247, 503, 15, 271, 143, 399, 79, 
                      335, 207, 463, 47, 303, 175, 431, 111, 367, 239, 495, 31, 
                      287, 159, 415, 95, 351, 223, 479, 63, 319, 191, 447, 127, 
                      383, 255, 511, 0, 64, 32, 96, 16, 80, 48, 112, 8, 72, 40, 
                      104, 24, 88, 56, 120, 4, 68, 36, 100, 20, 84, 52, 116, 
                      3, 131, 67, 195, 35, 163, 99, 227};

    /* Test reading all possible literal-lengths */
    for (int i = 0; i < 288; i++) { 
        reset_pos();
        
        int code = codes[i];
        
        unsigned char buf[2]; 
        buf[0] = code & 255; 
        buf[1] = code >> 8;

        ASSERT_EQ (read_chunk((char *) buf, HUFFMAN_FIXED), i); 
    }
    
    /* Test reading all possible fixed distances */

    /*
     * Mapping of distance literal -> code, again, as laid out in memory. 
     */
    int codes_dist[30] = {0, 16, 8, 24, 4, 20, 12, 28, 2, 18, 10, 26, 6, 22, 
                          14, 30, 1, 17, 9, 25, 5, 21, 13, 29, 3, 19, 11, 
                          27, 7, 23}; 

    for (int i = 0; i < 30; i++) { 
        reset_pos(); 

        int code = codes_dist[i];
        unsigned char buf[1];

        buf[0] = code; 
        ASSERT_EQ (read_chunk((char *) buf, HUFFMAN_FIXED_DISTS), i);
    }

}

int main(int argc, char **argv) { 
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}