#include "unity/unity.h"
#include "zlib.h"
#include "deflate.h" /* Needed for deflate_state, MIN_MATCH, etc. */

#include <string.h>
#include <stdint.h>

/* Wrapper provided in deflate.c for testing local function lm_init */
extern void test_lm_init(deflate_state *s);

void setUp(void) {
    /* Setup code here, or leave empty */
}

void tearDown(void) {
    /* Cleanup code here, or leave empty */
}

/* Helper to initialize a z_stream and return its internal deflate_state* */
static int init_stream_get_state(z_stream *strm, int level, deflate_state **out_s) {
    memset(strm, 0, sizeof(*strm));
    /* Use common parameters to ensure allocations happen */
    int ret = deflateInit2(strm, level, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY);
    if (ret != Z_OK) return ret;
    *out_s = (deflate_state *)strm->state;
    return Z_OK;
}

void test_lm_init_resets_state_fields_and_clears_hash(void) {
    z_stream strm;
    deflate_state *s = NULL;
    int ret = init_stream_get_state(&strm, 6, &s);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(s);
    TEST_ASSERT_NOT_EQUAL(0u, s->hash_size);
    TEST_ASSERT_NOT_NULL(s->head);

    /* Corrupt/alter fields to verify they are reset */
    s->window_size = 0;
    s->strstart = 12345;
    s->block_start = 999L;
    s->lookahead = 77;
    s->insert = 33;
    s->match_length = 42;
    s->prev_length = 17;
    s->match_available = 1;
    s->ins_h = 99;

    /* Make sure hash head table is non-zero everywhere */
    for (uInt i = 0; i < s->hash_size; i++) {
        s->head[i] = 1;
    }

    /* Call the wrapper to test lm_init */
    test_lm_init(s);

    /* window_size is set to 2 * w_size */
    TEST_ASSERT_EQUAL_UINT64((uint64_t)s->w_size * 2u, (uint64_t)s->window_size);

    /* Core fields are reset */
    TEST_ASSERT_EQUAL_UINT32(0u, s->strstart);
    TEST_ASSERT_EQUAL_INT32(0L, s->block_start);
    TEST_ASSERT_EQUAL_UINT32(0u, s->lookahead);
    TEST_ASSERT_EQUAL_UINT32(0u, s->insert);
    TEST_ASSERT_EQUAL_UINT32((unsigned)(MIN_MATCH - 1), s->match_length);
    TEST_ASSERT_EQUAL_UINT32((unsigned)(MIN_MATCH - 1), s->prev_length);
    TEST_ASSERT_EQUAL_UINT32(0u, s->match_available);
    TEST_ASSERT_EQUAL_UINT32(0u, s->ins_h);

    /* Hash head table cleared to NIL (0) */
    unsigned nonzero = 0;
    for (uInt i = 0; i < s->hash_size; i++) {
        if (s->head[i] != 0) {
            nonzero++;
            break; /* Early exit on first failure to reduce runtime */
        }
    }
    TEST_ASSERT_EQUAL_UINT(0u, nonzero);

    /* Cleanup */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

void test_lm_init_config_level0(void) {
    z_stream strm;
    deflate_state *s = NULL;
    int ret = init_stream_get_state(&strm, 0, &s);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(s);

    /* Set the desired level and poison config fields */
    s->level = 0;
    s->good_match = 1234;
    s->max_lazy_match = 4321;
    s->nice_match = 555;
    s->max_chain_length = 777;

    test_lm_init(s);

    /* From configuration_table[0]: {0,0,0,0,...} */
    TEST_ASSERT_EQUAL_UINT(0u, s->good_match);
    TEST_ASSERT_EQUAL_UINT(0u, s->max_lazy_match);
    TEST_ASSERT_EQUAL_UINT(0u, s->nice_match);
    TEST_ASSERT_EQUAL_UINT(0u, s->max_chain_length);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

void test_lm_init_config_level1(void) {
    z_stream strm;
    deflate_state *s = NULL;
    int ret = init_stream_get_state(&strm, 1, &s);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(s);

    s->level = 1;
    s->good_match = 0;
    s->max_lazy_match = 0;
    s->nice_match = 0;
    s->max_chain_length = 0;

    test_lm_init(s);

#ifdef FASTEST
    /* FASTEST configuration_table[1]: {4,4,8,4} */
    TEST_ASSERT_EQUAL_UINT(4u, s->good_match);
    TEST_ASSERT_EQUAL_UINT(4u, s->max_lazy_match);
    TEST_ASSERT_EQUAL_UINT(8u, s->nice_match);
    TEST_ASSERT_EQUAL_UINT(4u, s->max_chain_length);
#else
    /* Non-FASTEST configuration_table[1]: {4,4,8,4} as well */
    TEST_ASSERT_EQUAL_UINT(4u, s->good_match);
    TEST_ASSERT_EQUAL_UINT(4u, s->max_lazy_match);
    TEST_ASSERT_EQUAL_UINT(8u, s->nice_match);
    TEST_ASSERT_EQUAL_UINT(4u, s->max_chain_length);
#endif

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

#ifndef FASTEST
void test_lm_init_config_level6(void) {
    z_stream strm;
    deflate_state *s = NULL;
    int ret = init_stream_get_state(&strm, 6, &s);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(s);

    s->level = 6;
    s->good_match = 0;
    s->max_lazy_match = 0;
    s->nice_match = 0;
    s->max_chain_length = 0;

    test_lm_init(s);

    /* configuration_table[6]: {8, 16, 128, 128} */
    TEST_ASSERT_EQUAL_UINT(8u, s->good_match);
    TEST_ASSERT_EQUAL_UINT(16u, s->max_lazy_match);
    TEST_ASSERT_EQUAL_UINT(128u, s->nice_match);
    TEST_ASSERT_EQUAL_UINT(128u, s->max_chain_length);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

void test_lm_init_config_level9(void) {
    z_stream strm;
    deflate_state *s = NULL;
    int ret = init_stream_get_state(&strm, 9, &s);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(s);

    s->level = 9;
    s->good_match = 0;
    s->max_lazy_match = 0;
    s->nice_match = 0;
    s->max_chain_length = 0;

    test_lm_init(s);

    /* configuration_table[9]: {32, 258, 258, 4096} */
    TEST_ASSERT_EQUAL_UINT(32u, s->good_match);
    TEST_ASSERT_EQUAL_UINT(258u, s->max_lazy_match);
    TEST_ASSERT_EQUAL_UINT(258u, s->nice_match);
    TEST_ASSERT_EQUAL_UINT(4096u, s->max_chain_length);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}
#endif /* !FASTEST */

int main(void) {
    UNITY_BEGIN();

    RUN_TEST(test_lm_init_resets_state_fields_and_clears_hash);
    RUN_TEST(test_lm_init_config_level0);
    RUN_TEST(test_lm_init_config_level1);
#ifndef FASTEST
    RUN_TEST(test_lm_init_config_level6);
    RUN_TEST(test_lm_init_config_level9);
#endif

    return UNITY_END();
}