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

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

/* Helper calculators mirroring deflateBound() internal formulas */
static uLong calc_fixed(uLong n) {
    /* fixedlen = n + (n>>3) + (n>>8) + (n>>9) + 4 */
    return n + (n >> 3) + (n >> 8) + (n >> 9) + 4u;
}
static uLong calc_store(uLong n) {
    /* storelen = n + (n>>5) + (n>>7) + (n>>11) + 7 */
    return n + (n >> 5) + (n >> 7) + (n >> 11) + 7u;
}
static uLong calc_default(uLong n, uLong wraplen) {
    /* default tight bound (w_bits==15 && hash_bits==15):
       n + (n>>12) + (n>>14) + (n>>25) + 13 - 6 + wraplen */
    return n + (n >> 12) + (n >> 14) + (n >> 25) + 13u - 6u + wraplen;
}

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

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

/* 1) Invalid stream: NULL pointer */
void test_deflateBound_invalid_stream_null_ptr(void) {
    uLong src = 123456u;
    uLong expected = ((calc_fixed(src) > calc_store(src)) ? calc_fixed(src) : calc_store(src)) + 18u;
    uLong got = deflateBound(NULL, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);
}

/* 2) Invalid stream: zero-initialized z_stream (no allocators) */
void test_deflateBound_invalid_stream_uninitialized(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    uLong src = 98765u;
    uLong expected = ((calc_fixed(src) > calc_store(src)) ? calc_fixed(src) : calc_store(src)) + 18u;
    uLong got = deflateBound(&strm, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);
}

/* 3) Default stream, zlib wrapper, no preset dictionary */
void test_deflateBound_default_stream_no_dict(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit(&strm, Z_DEFAULT_COMPRESSION));
    /* After deflateInit with defaults: wrap=1 (zlib), w_bits=15, hash_bits=15, strstart==0 */
    uLong src = 100000u;
    uLong wraplen = 6u; /* zlib wrapper with no preset dictionary */
    uLong expected = calc_default(src, wraplen);
    uLong got = deflateBound(&strm, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

/* 4) Default stream with preset dictionary -> wraplen increases by 4 */
void test_deflateBound_default_stream_with_preset_dict(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit(&strm, Z_DEFAULT_COMPRESSION));

    /* Set a small preset dictionary; valid only in INIT_STATE with wrap==1 */
    const unsigned char dict[] = "preset-dictionary-bytes";
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateSetDictionary(&strm, dict, (uInt)sizeof(dict)));

    /* strstart > 0 now; wrap remains 1 -> wraplen = 6 + 4 = 10 */
    uLong src = 54321u;
    uLong wraplen = 10u;
    uLong expected = calc_default(src, wraplen);
    uLong got = deflateBound(&strm, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);

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

/* 5) Non-default memLevel (memLevel=1): expect conservative storelen + wraplen */
void test_deflateBound_nondefault_memLevel_storelen_path(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    /* windowBits=15, memLevel=1 -> hash_bits != 15, choose conservative bound
       w_bits (15) > hash_bits (8), level non-zero -> choose storelen */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit2(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED, 15, 1, Z_DEFAULT_STRATEGY));

    uLong src = 200000u;
    uLong wraplen = 6u; /* zlib wrapper, strstart==0 */
    uLong expected = calc_store(src) + wraplen;
    uLong got = deflateBound(&strm, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);

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

/* 6) Small window relative to memLevel (windowBits <= memLevel + 7): fixedlen + wraplen */
void test_deflateBound_small_window_fixedlen_path(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    /* windowBits=9, memLevel=8 -> w_bits (9) <= hash_bits (15), level non-zero -> fixedlen */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit2(&strm, 6, Z_DEFLATED, 9, 8, Z_DEFAULT_STRATEGY));

    uLong src = 65536u;
    uLong wraplen = 6u; /* zlib wrapper */
    uLong expected = calc_fixed(src) + wraplen;
    uLong got = deflateBound(&strm, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);

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

/* 7) Raw deflate (no wrapper): default parameters but wraplen = 0 */
void test_deflateBound_raw_deflate_no_wrapper(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    /* Negative windowBits -> raw deflate (wrap=0); w_bits will be 15; memLevel default 8 -> hash_bits=15 */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit2(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY));

    uLong src = 314159u;
    uLong wraplen = 0u; /* raw */
    uLong expected = calc_default(src, wraplen);
    uLong got = deflateBound(&strm, src);
    TEST_ASSERT_EQUAL_UINT64(expected, got);

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

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflateBound_invalid_stream_null_ptr);
    RUN_TEST(test_deflateBound_invalid_stream_uninitialized);
    RUN_TEST(test_deflateBound_default_stream_no_dict);
    RUN_TEST(test_deflateBound_default_stream_with_preset_dict);
    RUN_TEST(test_deflateBound_nondefault_memLevel_storelen_path);
    RUN_TEST(test_deflateBound_small_window_fixedlen_path);
    RUN_TEST(test_deflateBound_raw_deflate_no_wrapper);
    return UNITY_END();
}