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

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

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

/* Helper to allocate and check non-null */
static void* xmalloc(size_t n) {
    void* p = malloc(n);
    TEST_ASSERT_NOT_NULL(p);
    return p;
}

/* Test 1: Round-trip small text with typical compression level */
void test_compress2_roundtrip_small_text(void) {
    const char *msg = "The quick brown fox jumps over the lazy dog";
    const uLong srcLen = (uLong)strlen(msg);
    uLongf compCap = compressBound(srcLen);
    Bytef *comp = (Bytef*)xmalloc(compCap);
    uLongf compLen = compCap;

    int err = compress2(comp, &compLen, (const Bytef*)msg, srcLen, 6);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_TRUE(compLen > 0);
    TEST_ASSERT_TRUE(compLen <= compCap);

    Bytef *decomp = (Bytef*)xmalloc(srcLen ? srcLen : 1);
    uLongf decompLen = srcLen;
    err = uncompress(decomp, &decompLen, comp, compLen);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_EQUAL_UINT32(srcLen, decompLen);
    if (srcLen > 0) {
        TEST_ASSERT_EQUAL_UINT8_ARRAY((const uint8_t*)msg, decomp, srcLen);
    }

    free(comp);
    free(decomp);
}

/* Test 2: Zero-length input should succeed and decompress to zero bytes */
void test_compress2_zero_input(void) {
    const Bytef *src = (const Bytef*)""; /* non-null pointer */
    const uLong srcLen = 0;
    uLongf compCap = compressBound(srcLen);
    if (compCap < 16) compCap = 16; /* ensure some room */
    Bytef *comp = (Bytef*)xmalloc(compCap);
    uLongf compLen = compCap;

    int err = compress2(comp, &compLen, src, srcLen, Z_BEST_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_TRUE(compLen > 0);
    TEST_ASSERT_TRUE(compLen <= compCap);

    Bytef *decomp = (Bytef*)xmalloc(1); /* capacity 1, expect output size 0 */
    uLongf decompLen = 0;
    err = uncompress(decomp, &decompLen, comp, compLen);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_EQUAL_UINT32(0, decompLen);

    free(comp);
    free(decomp);
}

/* Test 3: Output buffer too small should return Z_BUF_ERROR and update destLen */
void test_compress2_buffer_too_small(void) {
    /* Prepare some input */
    const uLong srcLen = 1000;
    Bytef *src = (Bytef*)xmalloc(srcLen);
    memset(src, 'A', srcLen);

    Bytef smallBuf[1];
    uLongf smallLen = sizeof(smallBuf);

    int err = compress2(smallBuf, &smallLen, src, srcLen, Z_BEST_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, err);
    /* destLen should be the number of bytes produced before running out of space */
    TEST_ASSERT_TRUE(smallLen <= sizeof(smallBuf));

    free(src);
}

/* Test 4: Zero capacity buffer should return Z_BUF_ERROR and leave destLen 0 */
void test_compress2_zero_capacity(void) {
    const uLong srcLen = 256;
    Bytef *src = (Bytef*)xmalloc(srcLen);
    for (uLong i = 0; i < srcLen; ++i) src[i] = (Bytef)(i & 0xFF);

    Bytef *dest = NULL; /* no space */
    uLongf destLen = 0;

    int err = compress2(dest, &destLen, src, srcLen, 6);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, err);
    TEST_ASSERT_EQUAL_UINT32(0, destLen);

    free(src);
}

/* Test 5: Invalid compression level should return Z_STREAM_ERROR and set destLen to 0 */
void test_compress2_invalid_level(void) {
    const char *msg = "invalid level test data";
    const uLong srcLen = (uLong)strlen(msg);
    Bytef destBuf[128];
    uLongf destLen = sizeof(destBuf); /* non-zero to verify it becomes 0 on error */

    int err = compress2(destBuf, &destLen, (const Bytef*)msg, srcLen, 99); /* invalid level */
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, err);
    TEST_ASSERT_EQUAL_UINT32(0, destLen);
}

/* Test 6: Large input round-trip at levels 0 and 9; level 9 should compress better */
void test_compress2_large_input_levels(void) {
    const uLong srcLen = 200000; /* 200 KB */
    Bytef *src = (Bytef*)xmalloc(srcLen);
    memset(src, 'A', srcLen); /* highly compressible */

    /* Level 0 (store) */
    uLongf cap0 = compressBound(srcLen);
    Bytef *comp0 = (Bytef*)xmalloc(cap0);
    uLongf len0 = cap0;
    int err = compress2(comp0, &len0, src, srcLen, 0);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_TRUE(len0 > 0);
    TEST_ASSERT_TRUE(len0 <= cap0);

    /* Decompress and verify */
    Bytef *decomp0 = (Bytef*)xmalloc(srcLen);
    uLongf dlen0 = srcLen;
    err = uncompress(decomp0, &dlen0, comp0, len0);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_EQUAL_UINT32(srcLen, dlen0);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(src, decomp0, srcLen);

    /* Level 9 (best) */
    uLongf cap9 = compressBound(srcLen);
    Bytef *comp9 = (Bytef*)xmalloc(cap9);
    uLongf len9 = cap9;
    err = compress2(comp9, &len9, src, srcLen, 9);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_TRUE(len9 > 0);
    TEST_ASSERT_TRUE(len9 <= cap9);

    /* Decompress and verify */
    Bytef *decomp9 = (Bytef*)xmalloc(srcLen);
    uLongf dlen9 = srcLen;
    err = uncompress(decomp9, &dlen9, comp9, len9);
    TEST_ASSERT_EQUAL_INT(Z_OK, err);
    TEST_ASSERT_EQUAL_UINT32(srcLen, dlen9);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(src, decomp9, srcLen);

    /* Expect better compression at level 9 than at level 0 for repetitive data */
    TEST_ASSERT_TRUE(len9 <= len0);

    free(src);
    free(comp0);
    free(decomp0);
    free(comp9);
    free(decomp9);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_compress2_roundtrip_small_text);
    RUN_TEST(test_compress2_zero_input);
    RUN_TEST(test_compress2_buffer_too_small);
    RUN_TEST(test_compress2_zero_capacity);
    RUN_TEST(test_compress2_invalid_level);
    RUN_TEST(test_compress2_large_input_levels);
    return UNITY_END();
}