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

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

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

static void init_stream(z_stream *strm) {
    memset(strm, 0, sizeof(*strm));
    int ret = deflateInit(strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

/* Test: NULL stream pointer returns Z_STREAM_ERROR */
void test_deflateParams_null_stream_returns_error(void) {
    int ret = deflateParams(NULL, 1, Z_DEFAULT_STRATEGY);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: Uninitialized stream (zalloc/zfree are NULL) returns Z_STREAM_ERROR */
void test_deflateParams_uninitialized_stream_returns_error(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateParams(&strm, 1, Z_DEFAULT_STRATEGY);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: Invalid level values return Z_STREAM_ERROR */
void test_deflateParams_invalid_level_returns_error(void) {
    z_stream strm;
    init_stream(&strm);

    int ret;
    ret = deflateParams(&strm, 10, Z_DEFAULT_STRATEGY); /* > 9 */
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    ret = deflateParams(&strm, -2, Z_DEFAULT_STRATEGY); /* < -1, note -1 is Z_DEFAULT_COMPRESSION */
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    deflateEnd(&strm);
}

/* Test: Invalid strategy values return Z_STREAM_ERROR */
void test_deflateParams_invalid_strategy_returns_error(void) {
    z_stream strm;
    init_stream(&strm);

    int ret = deflateParams(&strm, 1, Z_FIXED + 1);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    deflateEnd(&strm);
}

/* Test: Z_DEFAULT_COMPRESSION is accepted and returns Z_OK */
void test_deflateParams_accepts_default_compression_constant(void) {
    z_stream strm;
    init_stream(&strm);

    int ret = deflateParams(&strm, Z_DEFAULT_COMPRESSION, Z_DEFAULT_STRATEGY);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    deflateEnd(&strm);
}

/* Test: On a fresh stream (last_flush == -2), deflateParams does not flush and returns Z_OK even if input is pending */
void test_deflateParams_fresh_stream_with_pending_input_returns_ok(void) {
    z_stream strm;
    init_stream(&strm);

    unsigned char inbuf[128];
    for (size_t i = 0; i < sizeof(inbuf); i++) inbuf[i] = (unsigned char)(i & 0xFF);
    unsigned char outbuf[1]; /* non-NULL pointer is required, size doesn't matter here */

    strm.next_in = inbuf;
    strm.avail_in = (uInt)sizeof(inbuf);
    strm.next_out = outbuf;
    strm.avail_out = (uInt)sizeof(outbuf);

    /* Change strategy to force the (strategy != s->strategy) condition, but since last_flush == -2, no flush happens */
    int ret = deflateParams(&strm, 6 /* keep level same as default resolved level */, Z_HUFFMAN_ONLY);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    deflateEnd(&strm);
}

/* Helper: set last_flush != -2 without consuming input by calling deflate with avail_out == 0 */
static void set_last_flush_non_negative(z_stream *strm) {
    /* Provide non-NULL next_out with zero avail_out to set last_flush via deflate() */
    unsigned char dummy;
    strm->next_out = &dummy;
    strm->avail_out = 0; /* forces deflate() to return Z_BUF_ERROR, but sets last_flush */
    strm->next_in = Z_NULL;
    strm->avail_in = 0;
    int ret = deflate(strm, Z_NO_FLUSH);
    /* We expect Z_BUF_ERROR here due to zero avail_out, that's fine */
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, ret);
}

/* Test: When last_flush != -2 and input pending, deflateParams will attempt a flush and return Z_BUF_ERROR */
void test_deflateParams_triggers_flush_and_returns_buf_error_with_pending_input(void) {
    z_stream strm;
    init_stream(&strm);

    /* First, set last_flush to a non -2 value */
    set_last_flush_non_negative(&strm);

    /* Now set pending input and zero output space */
    unsigned char inbuf[256];
    for (size_t i = 0; i < sizeof(inbuf); i++) inbuf[i] = (unsigned char)('A' + (i % 26));
    unsigned char dummy_out;
    strm.next_in = inbuf;
    strm.avail_in = (uInt)sizeof(inbuf);
    strm.next_out = &dummy_out;
    strm.avail_out = 0; /* ensure deflate(Z_BLOCK) can't output */

    /* Change strategy to force the flush path */
    int ret = deflateParams(&strm, 6 /* same level func */, Z_HUFFMAN_ONLY);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, ret);

    deflateEnd(&strm);
}

/* Test: When last_flush != -2 and no input and no pending block, deflateParams returns Z_OK even if internal deflate(Z_BLOCK) couldn't write */
void test_deflateParams_flush_path_no_pending_returns_ok(void) {
    z_stream strm;
    init_stream(&strm);

    /* Set last_flush from -2 to 0 using deflate with zero avail_out */
    set_last_flush_non_negative(&strm);

    /* No input pending, keep zero output space, change level to alter compression function */
    unsigned char dummy_out;
    strm.next_in = Z_NULL;
    strm.avail_in = 0;
    strm.next_out = &dummy_out;
    strm.avail_out = 0;

    /* Change level from default resolved level (6) to 1 to change func (slow -> fast) */
    int ret = deflateParams(&strm, 1, Z_DEFAULT_STRATEGY);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    deflateEnd(&strm);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflateParams_null_stream_returns_error);
    RUN_TEST(test_deflateParams_uninitialized_stream_returns_error);
    RUN_TEST(test_deflateParams_invalid_level_returns_error);
    RUN_TEST(test_deflateParams_invalid_strategy_returns_error);
    RUN_TEST(test_deflateParams_accepts_default_compression_constant);
    RUN_TEST(test_deflateParams_fresh_stream_with_pending_input_returns_ok);
    RUN_TEST(test_deflateParams_triggers_flush_and_returns_buf_error_with_pending_input);
    RUN_TEST(test_deflateParams_flush_path_no_pending_returns_ok);
    return UNITY_END();
}