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

/* Wrapper provided in deflate.c for testing local function */
extern int test_deflateStateCheck(z_streamp strm);

void setUp(void) {
    /* no setup needed */
}

void tearDown(void) {
    /* no teardown needed */
}

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);
}

void test_deflateStateCheck_NULL_stream_returns_1(void) {
    int rc = test_deflateStateCheck(NULL);
    TEST_ASSERT_EQUAL_INT(1, rc);
}

void test_deflateStateCheck_uninitialized_stream_returns_1(void) {
    z_stream s;
    memset(&s, 0, sizeof(s));
    /* zalloc and zfree are NULL; state is NULL */
    int rc = test_deflateStateCheck(&s);
    TEST_ASSERT_EQUAL_INT(1, rc);
}

void test_deflateStateCheck_after_deflateInit_is_valid(void) {
    z_stream strm;
    init_stream(&strm);

    int rc = test_deflateStateCheck(&strm);
    TEST_ASSERT_EQUAL_INT(0, rc);

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

void test_deflateStateCheck_in_BUSY_STATE_after_header_is_valid(void) {
    z_stream strm;
    init_stream(&strm);

    unsigned char out[128];
    strm.next_out = out;
    strm.avail_out = (uInt)sizeof(out);
    /* No input needed to emit zlib header and enter BUSY_STATE */
    int ret = deflate(&strm, Z_NO_FLUSH);
    TEST_ASSERT(ret == Z_OK || ret == Z_BUF_ERROR /* overly small out buf case */);

    int rc = test_deflateStateCheck(&strm);
    TEST_ASSERT_EQUAL_INT(0, rc);

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

void test_deflateStateCheck_in_FINISH_STATE_is_valid(void) {
    z_stream strm;
    init_stream(&strm);

    const unsigned char in[] = "Hello, zlib!";
    unsigned char out[1024];

    strm.next_in = (Bytef *)in;
    strm.avail_in = (uInt)sizeof(in) - 1; /* exclude terminating NUL */
    strm.next_out = out;
    strm.avail_out = (uInt)sizeof(out);

    int ret;
    do {
        ret = deflate(&strm, Z_FINISH);
        /* keep calling until stream end */
    } while (ret == Z_OK);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, ret);

    int rc = test_deflateStateCheck(&strm);
    TEST_ASSERT_EQUAL_INT(0, rc);

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

void test_deflateStateCheck_with_state_from_different_stream_returns_1(void) {
    z_stream owner;
    init_stream(&owner);

    /* Build a second z_stream that points to the first stream's state. */
    z_stream imposter;
    memset(&imposter, 0, sizeof(imposter));
    /* Satisfy allocator checks so we get past the first condition */
    imposter.zalloc = owner.zalloc;
    imposter.zfree = owner.zfree;
    imposter.opaque = owner.opaque;
    /* Illegally point at someone else's state */
    imposter.state = owner.state;

    int rc = test_deflateStateCheck(&imposter);
    TEST_ASSERT_EQUAL_INT(1, rc);

    /* Clean up only the legitimate owner */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&owner));
}

void test_deflateStateCheck_after_deflateEnd_returns_1(void) {
    z_stream strm;
    init_stream(&strm);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));

    /* zalloc/zfree are still set, but state is NULL -> should be invalid */
    int rc = test_deflateStateCheck(&strm);
    TEST_ASSERT_EQUAL_INT(1, rc);
}

int main(void) {
    UNITY_BEGIN();

    RUN_TEST(test_deflateStateCheck_NULL_stream_returns_1);
    RUN_TEST(test_deflateStateCheck_uninitialized_stream_returns_1);
    RUN_TEST(test_deflateStateCheck_after_deflateInit_is_valid);
    RUN_TEST(test_deflateStateCheck_in_BUSY_STATE_after_header_is_valid);
    RUN_TEST(test_deflateStateCheck_in_FINISH_STATE_is_valid);
    RUN_TEST(test_deflateStateCheck_with_state_from_different_stream_returns_1);
    RUN_TEST(test_deflateStateCheck_after_deflateEnd_returns_1);

    return UNITY_END();
}