#include "unity/unity.h"
#include "zlib.h"
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

/* Globals for custom allocator tests */
static int g_alloc_calls;
static int g_free_calls;
static int g_fail_on_alloc_call; /* 0 means no failure */

static voidpf my_zalloc(voidpf opaque, uInt items, uInt size) {
    (void)opaque;
    g_alloc_calls++;
    if (g_fail_on_alloc_call > 0 && g_alloc_calls == g_fail_on_alloc_call) {
        return Z_NULL; /* simulate allocation failure */
    }
    size_t total = (size_t)items * (size_t)size;
    return malloc(total ? total : 1);
}

static void my_zfree(voidpf opaque, voidpf address) {
    (void)opaque;
    g_free_calls++;
    free(address);
}

void setUp(void) {
    g_alloc_calls = 0;
    g_free_calls = 0;
    g_fail_on_alloc_call = 0;
}

void tearDown(void) {
    /* nothing */
}

/* Helper to safely end a stream if it was initialized */
static void end_if_inited(z_stream *strm) {
    if (strm && strm->state != Z_NULL) {
        deflateEnd(strm);
    }
}

void test_deflateInit2__returns_Z_OK_and_allocates_state_default(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit2_(&strm,
                            Z_DEFAULT_COMPRESSION,
                            Z_DEFLATED,
                            15,
                            8,
                            Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION,
                            (int)sizeof(strm));

    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(strm.state);
    TEST_ASSERT_NOT_NULL(strm.zalloc);
    TEST_ASSERT_NOT_NULL(strm.zfree);

    ret = deflateEnd(&strm);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NULL(strm.state);
}

void test_deflateInit2__null_stream_returns_Z_STREAM_ERROR(void) {
    int ret = deflateInit2_(Z_NULL,
                            Z_DEFAULT_COMPRESSION,
                            Z_DEFLATED,
                            15,
                            8,
                            Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION,
                            (int)sizeof(z_stream));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

void test_deflateInit2__wrong_version_returns_Z_VERSION_ERROR_and_no_state(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit2_(&strm,
                            Z_DEFAULT_COMPRESSION,
                            Z_DEFLATED,
                            15,
                            8,
                            Z_DEFAULT_STRATEGY,
                            "0", /* wrong major version */
                            (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_VERSION_ERROR, ret);
    TEST_ASSERT_NULL(strm.state);
}

void test_deflateInit2__wrong_size_returns_Z_VERSION_ERROR(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit2_(&strm,
                            Z_DEFAULT_COMPRESSION,
                            Z_DEFLATED,
                            15,
                            8,
                            Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION,
                            (int)sizeof(strm) - 1);
    TEST_ASSERT_EQUAL_INT(Z_VERSION_ERROR, ret);
}

void test_deflateInit2__invalid_memLevel_low_high(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 15, 0, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
    TEST_ASSERT_NULL(strm.state);

    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, 6, Z_DEFLATED, 15, MAX_MEM_LEVEL + 1, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
    TEST_ASSERT_NULL(strm.state);
}

void test_deflateInit2__invalid_method(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2_(&strm, 6, 0 /* not Z_DEFLATED */, 15, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

void test_deflateInit2__invalid_windowBits_values(void) {
    z_stream strm;

    /* windowBits too small */
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 7, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* windowBits = 16 -> wrap becomes 2, internal windowBits 0 -> invalid */
    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, 6, Z_DEFLATED, 16, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* windowBits < -15 invalid */
    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, 6, Z_DEFLATED, -16, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* windowBits == -8 invalid (wrap 0 with 8) */
    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, 6, Z_DEFLATED, -8, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* Negative but valid (raw); -9 should succeed */
    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, 6, Z_DEFLATED, -9, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    end_if_inited(&strm);
}

void test_deflateInit2__invalid_strategy_range(void) {
    z_stream strm; memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 15, 8, Z_FIXED + 1,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

void test_deflateInit2__level_range_and_default_acceptance(void) {
    z_stream strm;

    /* Level too high */
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2_(&strm, 10, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* Level invalid negative (not default) */
    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, -2, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* Default level (-1) should be accepted */
    memset(&strm, 0, sizeof(strm));
    ret = deflateInit2_(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    end_if_inited(&strm);
}

void test_deflateInit2__windowBits8_wrap1_ok(void) {
    z_stream strm; memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 8, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    end_if_inited(&strm);
}

void test_deflateInit2__wrapper_behavior_via_deflateSetHeader(void) {
    gz_header hdr;
    memset(&hdr, 0, sizeof(hdr));
    int ret;

    /* Gzip wrapper: windowBits > 15 (e.g., 31) */
    z_stream strm_gz; memset(&strm_gz, 0, sizeof(strm_gz));
    ret = deflateInit2_(&strm_gz, 6, Z_DEFLATED, 31, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm_gz));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    int shret = deflateSetHeader(&strm_gz, &hdr);
    TEST_ASSERT_EQUAL_INT(Z_OK, shret);
    end_if_inited(&strm_gz);

    /* Raw deflate: negative windowBits (e.g., -15) */
    z_stream strm_raw; memset(&strm_raw, 0, sizeof(strm_raw));
    ret = deflateInit2_(&strm_raw, 6, Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm_raw));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    shret = deflateSetHeader(&strm_raw, &hdr);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, shret);
    end_if_inited(&strm_raw);

    /* Zlib wrapper: standard 15 should not accept gzip header */
    z_stream strm_z; memset(&strm_z, 0, sizeof(strm_z));
    ret = deflateInit2_(&strm_z, 6, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                        ZLIB_VERSION, (int)sizeof(strm_z));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    shret = deflateSetHeader(&strm_z, &hdr);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, shret);
    end_if_inited(&strm_z);
}

void test_deflateInit2__uses_custom_allocators(void) {
    z_stream strm; memset(&strm, 0, sizeof(strm));
    strm.zalloc = my_zalloc;
    strm.zfree = my_zfree;
    strm.opaque = (voidpf)0xDEADBEEF;

    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_TRUE(g_alloc_calls > 0);

    ret = deflateEnd(&strm);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_TRUE(g_free_calls > 0);
}

void test_deflateInit2__allocation_failure_returns_Z_MEM_ERROR_and_cleans_up(void) {
    z_stream strm; memset(&strm, 0, sizeof(strm));
    strm.zalloc = my_zalloc;
    strm.zfree = my_zfree;
    /* Fail on the second allocation (first is deflate_state, second is window) */
    g_fail_on_alloc_call = 2;

    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_MEM_ERROR, ret);
    /* deflateInit2_ should have cleaned up and nulled the state */
    TEST_ASSERT_NULL(strm.state);
    TEST_ASSERT_TRUE(g_free_calls >= 1);
}

void test_deflateInit2__deflateReset_succeeds_after_init(void) {
    z_stream strm; memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2_(&strm, 6, Z_DEFLATED, 15, 8, Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION, (int)sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    int r2 = deflateReset(&strm);
    TEST_ASSERT_EQUAL_INT(Z_OK, r2);

    end_if_inited(&strm);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflateInit2__returns_Z_OK_and_allocates_state_default);
    RUN_TEST(test_deflateInit2__null_stream_returns_Z_STREAM_ERROR);
    RUN_TEST(test_deflateInit2__wrong_version_returns_Z_VERSION_ERROR_and_no_state);
    RUN_TEST(test_deflateInit2__wrong_size_returns_Z_VERSION_ERROR);
    RUN_TEST(test_deflateInit2__invalid_memLevel_low_high);
    RUN_TEST(test_deflateInit2__invalid_method);
    RUN_TEST(test_deflateInit2__invalid_windowBits_values);
    RUN_TEST(test_deflateInit2__invalid_strategy_range);
    RUN_TEST(test_deflateInit2__level_range_and_default_acceptance);
    RUN_TEST(test_deflateInit2__windowBits8_wrap1_ok);
    RUN_TEST(test_deflateInit2__wrapper_behavior_via_deflateSetHeader);
    RUN_TEST(test_deflateInit2__uses_custom_allocators);
    RUN_TEST(test_deflateInit2__allocation_failure_returns_Z_MEM_ERROR_and_cleans_up);
    RUN_TEST(test_deflateInit2__deflateReset_succeeds_after_init);
    return UNITY_END();
}