#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 */
}

/* Test: Passing NULL should return Z_STREAM_ERROR */
void test_deflateEnd_null_stream_returns_stream_error(void) {
    int ret = deflateEnd(NULL);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: Uninitialized z_stream (no allocators/state) should return Z_STREAM_ERROR */
void test_deflateEnd_uninitialized_stream_returns_stream_error(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateEnd(&strm);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: After deflateInit, deflateEnd should return Z_OK and clear state */
void test_deflateEnd_after_init_returns_ok_and_clears_state(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_NOT_NULL(strm.state);

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

/* Test: After header is written (BUSY_STATE), deflateEnd should return Z_DATA_ERROR */
void test_deflateEnd_after_header_written_returns_data_error(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

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

    unsigned char outbuf[128];
    strm.next_out = outbuf;
    strm.avail_out = (uInt)sizeof(outbuf);

    /* No input, but calling deflate once will write the zlib header and set BUSY_STATE */
    ret = deflate(&strm, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    /* Now end the stream without finishing: expect Z_DATA_ERROR */
    ret = deflateEnd(&strm);
    TEST_ASSERT_EQUAL_INT(Z_DATA_ERROR, ret);
    TEST_ASSERT_NULL(strm.state);
}

/* Test: After finishing compression (Z_STREAM_END), deflateEnd should return Z_OK */
void test_deflateEnd_after_finish_returns_ok(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

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

    const char *msg = "hello world, this is a small test payload";
    unsigned char inbuf[128];
    size_t inlen = strlen(msg);
    TEST_ASSERT(inlen < sizeof(inbuf));
    memcpy(inbuf, msg, inlen);

    unsigned char outbuf[512];
    memset(outbuf, 0, sizeof(outbuf));

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

    ret = deflate(&strm, Z_FINISH);
    /* With ample output buffer, we should finish in one call, but allow Z_OK and loop if needed */
    while (ret == Z_OK) {
        if (strm.avail_out == 0) {
            /* Provide more output room if somehow exhausted */
            /* In this test, buffer is large enough; this path shouldn't happen */
            TEST_FAIL_MESSAGE("Unexpectedly ran out of output buffer");
        }
        ret = deflate(&strm, Z_FINISH);
    }
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, ret);

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

/* Test: Calling deflateEnd twice: second call should return Z_STREAM_ERROR */
void test_deflateEnd_called_twice_second_call_errors(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

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

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

    /* Second call: state is NULL, so this should be a stream error */
    ret = deflateEnd(&strm);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflateEnd_null_stream_returns_stream_error);
    RUN_TEST(test_deflateEnd_uninitialized_stream_returns_stream_error);
    RUN_TEST(test_deflateEnd_after_init_returns_ok_and_clears_state);
    RUN_TEST(test_deflateEnd_after_header_written_returns_data_error);
    RUN_TEST(test_deflateEnd_after_finish_returns_ok);
    RUN_TEST(test_deflateEnd_called_twice_second_call_errors);
    return UNITY_END();
}