#include "unity/unity.h"
#include "zlib.h"

#include <string.h>
#include <stdlib.h>

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

/* Helper: initialize a fresh deflater stream */
static void init_deflater(z_stream *strm) {
    memset(strm, 0, sizeof(*strm));
    int ret = deflateInit(strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

/* Helper: get current pending and bit count */
static void get_pending_bits(z_stream *strm, unsigned *pending, int *bits) {
    int ret = deflatePending(strm, pending, bits);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

/* Test: invalid stream returns Z_STREAM_ERROR */
void test_deflatePrime_invalid_stream_returns_stream_error(void) {
    z_stream s;
    memset(&s, 0, sizeof(s)); /* invalid: zalloc/zfree are NULL */
    int ret = deflatePrime(&s, 1, 1);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: negative bits => Z_BUF_ERROR and no state change */
void test_deflatePrime_negative_bits_returns_buf_error_no_change(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending_before = 0;
    int bits_before = -1;
    get_pending_bits(&s, &pending_before, &bits_before);

    int ret = deflatePrime(&s, -1, 0);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, ret);

    unsigned pending_after = 0;
    int bits_after = -1;
    get_pending_bits(&s, &pending_after, &bits_after);

    TEST_ASSERT_EQUAL_UINT(pending_before, pending_after);
    TEST_ASSERT_EQUAL_INT(bits_before, bits_after);

    deflateEnd(&s);
}

/* Test: bits > 16 => Z_BUF_ERROR and no state change */
void test_deflatePrime_bits_exceed_16_returns_buf_error_no_change(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending_before = 0;
    int bits_before = -1;
    get_pending_bits(&s, &pending_before, &bits_before);

    int ret = deflatePrime(&s, 17, 0x1FFFF);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, ret);

    unsigned pending_after = 0;
    int bits_after = -1;
    get_pending_bits(&s, &pending_after, &bits_after);

    TEST_ASSERT_EQUAL_UINT(pending_before, pending_after);
    TEST_ASSERT_EQUAL_INT(bits_before, bits_after);

    deflateEnd(&s);
}

/* Test: bits == 0 is a no-op (Z_OK, no change) */
void test_deflatePrime_zero_bits_no_change(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending_before = 0;
    int bits_before = -1;
    get_pending_bits(&s, &pending_before, &bits_before);

    int ret = deflatePrime(&s, 0, 0xABCD);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    unsigned pending_after = 0;
    int bits_after = -1;
    get_pending_bits(&s, &pending_after, &bits_after);

    TEST_ASSERT_EQUAL_UINT(pending_before, pending_after);
    TEST_ASSERT_EQUAL_INT(bits_before, bits_after);

    deflateEnd(&s);
}

/* Test: accumulate less than one byte: pending unchanged, bits increased */
void test_deflatePrime_accumulate_less_than_byte(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending = 0;
    int bits = -1;
    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(0, pending);
    TEST_ASSERT_EQUAL_INT(0, bits);

    int ret = deflatePrime(&s, 3, 0x5); /* 101b */
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(0, pending); /* still no full byte flushed */
    TEST_ASSERT_EQUAL_INT(3, bits);     /* 3 bits in buffer */

    deflateEnd(&s);
}

/* Test: cross 16-bit boundary over two calls -> pending += 2, bits == 0 */
void test_deflatePrime_cross_boundary_flushes_two_bytes(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending = 0;
    int bits = -1;
    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(0, pending);
    TEST_ASSERT_EQUAL_INT(0, bits);

    int ret = deflatePrime(&s, 3, 0x5);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    ret = deflatePrime(&s, 13, 0x1A2B); /* total now 16 bits */
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(2, pending); /* 16 bits flushed as two bytes */
    TEST_ASSERT_EQUAL_INT(0, bits);     /* bit buffer reset */

    deflateEnd(&s);
}

/* Test: exactly 16 bits in one call => pending += 2, bits == 0 */
void test_deflatePrime_sixteen_bits_flushes_two_bytes(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending = 0;
    int bits = -1;
    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(0, pending);
    TEST_ASSERT_EQUAL_INT(0, bits);

    int ret = deflatePrime(&s, 16, 0xBEEF);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(2, pending);
    TEST_ASSERT_EQUAL_INT(0, bits);

    deflateEnd(&s);
}

/* Test: multiple 16-bit inserts accumulate pending bytes */
void test_deflatePrime_multiple_16bit_calls_accumulate_pending(void) {
    z_stream s;
    init_deflater(&s);

    unsigned pending = 0;
    int bits = -1;

    int ret = deflatePrime(&s, 16, 0x1111);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(2, pending);
    TEST_ASSERT_EQUAL_INT(0, bits);

    ret = deflatePrime(&s, 16, 0x2222);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    get_pending_bits(&s, &pending, &bits);
    TEST_ASSERT_EQUAL_UINT(4, pending);
    TEST_ASSERT_EQUAL_INT(0, bits);

    deflateEnd(&s);
}

/* Test: using deflatePrime after deflateEnd returns Z_STREAM_ERROR */
void test_deflatePrime_after_deflateEnd_returns_stream_error(void) {
    z_stream s;
    init_deflater(&s);
    int end_ret = deflateEnd(&s);
    TEST_ASSERT_EQUAL_INT(Z_OK, end_ret);

    int ret = deflatePrime(&s, 8, 0xAA);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Unity main */
int main(void) {
    UNITY_BEGIN();

    RUN_TEST(test_deflatePrime_invalid_stream_returns_stream_error);
    RUN_TEST(test_deflatePrime_negative_bits_returns_buf_error_no_change);
    RUN_TEST(test_deflatePrime_bits_exceed_16_returns_buf_error_no_change);
    RUN_TEST(test_deflatePrime_zero_bits_no_change);
    RUN_TEST(test_deflatePrime_accumulate_less_than_byte);
    RUN_TEST(test_deflatePrime_cross_boundary_flushes_two_bytes);
    RUN_TEST(test_deflatePrime_sixteen_bits_flushes_two_bytes);
    RUN_TEST(test_deflatePrime_multiple_16bit_calls_accumulate_pending);
    RUN_TEST(test_deflatePrime_after_deflateEnd_returns_stream_error);

    return UNITY_END();
}