#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: compress input entirely with deflate(), single-shot finish.
 * Returns Z_OK or Z_STREAM_END from deflate loop; sets *out_len to produced bytes.
 */
static int compress_with_deflate_single(const Bytef *in, uLong in_len,
                                        Bytef *out, uLong out_cap,
                                        uLong *out_len) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    if (ret != Z_OK) return ret;

    strm.next_in = (Bytef *)in;
    strm.avail_in = (uInt)in_len;
    strm.next_out = out;
    strm.avail_out = (uInt)out_cap;

    int dret;
    do {
        dret = deflate(&strm, Z_FINISH);
        if (dret == Z_STREAM_ERROR) break;
        if (dret == Z_BUF_ERROR) {
            /* Need more output space */
            if (strm.avail_out == 0) {
                deflateEnd(&strm);
                return Z_BUF_ERROR;
            }
            /* Otherwise, no progress (shouldn't happen here) */
            break;
        }
    } while (dret != Z_STREAM_END);

    *out_len = strm.total_out;
    int end_ret = deflateEnd(&strm);
    (void)end_ret; /* ignore, we validate via dret and produced data */
    return dret;
}

/* Helper: streaming compression with two phases (NO_FLUSH then FINISH),
 * while collecting output in an aggregate buffer.
 */
static int compress_with_deflate_stream_two_phase(const Bytef *in, uLong in_len,
                                                  Bytef *out, uLong out_cap,
                                                  uLong *out_len) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));

    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    if (ret != Z_OK) return ret;

    uLong produced = 0;

    /* Phase 1: first half with Z_NO_FLUSH */
    uLong half = in_len / 2;
    strm.next_in = (Bytef *)in;
    strm.avail_in = (uInt)half;
    strm.next_out = out + produced;
    strm.avail_out = (uInt)(out_cap - produced);

    int dret = deflate(&strm, Z_NO_FLUSH);
    if (dret != Z_OK && dret != Z_BUF_ERROR) {
        deflateEnd(&strm);
        return dret;
    }
    produced = strm.total_out;

    /* Phase 2: remaining input with Z_FINISH, loop to drain output */
    strm.next_in = (Bytef *)(in + half);
    strm.avail_in = (uInt)(in_len - half);

    do {
        strm.next_out = out + produced;
        strm.avail_out = (uInt)(out_cap - produced);
        dret = deflate(&strm, Z_FINISH);
        produced = strm.total_out;
        if (dret == Z_BUF_ERROR) {
            if (strm.avail_out == 0) {
                deflateEnd(&strm);
                return Z_BUF_ERROR;
            }
            /* If no progress but have room, continue */
        }
    } while (dret != Z_STREAM_END);

    *out_len = produced;
    int end_ret = deflateEnd(&strm);
    (void)end_ret;
    return dret;
}

/* Test 1: NULL/invalid stream yields Z_STREAM_ERROR */
void test_deflate_null_stream_error(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm)); /* invalid: no allocators set */
    int ret = deflate(&strm, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test 2: invalid flush value (> Z_BLOCK) yields Z_STREAM_ERROR */
void test_deflate_invalid_flush_value(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[16];
    strm.next_out = outbuf;
    strm.avail_out = sizeof(outbuf);

    ret = deflate(&strm, 999); /* invalid flush */
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

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

/* Test 3: next_out is NULL yields Z_STREAM_ERROR */
void test_deflate_null_next_out(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    strm.next_out = Z_NULL;
    strm.avail_out = 1; /* non-zero, to ensure we hit next_out NULL path */
    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    ret = deflate(&strm, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

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

/* Test 4: avail_out == 0 yields Z_BUF_ERROR */
void test_deflate_zero_avail_out(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[1];
    strm.next_out = outbuf;
    strm.avail_out = 0; /* provoke Z_BUF_ERROR */
    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    ret = deflate(&strm, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, ret);

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

/* Test 5: simple round-trip compression using deflate + uncompress */
void test_deflate_roundtrip_small(void) {
    const char *msg = "hello world";
    const uLong in_len = (uLong)strlen(msg);

    unsigned char comp[128];
    uLong comp_len = 0;
    int ret = compress_with_deflate_single((const Bytef *)msg, in_len, comp, sizeof(comp), &comp_len);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, ret);
    TEST_ASSERT_TRUE(comp_len > 0);

    unsigned char decomp[128];
    uLongf decomp_len = sizeof(decomp);
    int ur = uncompress(decomp, &decomp_len, comp, comp_len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ur);
    TEST_ASSERT_EQUAL_UINT(in_len, decomp_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(msg, decomp, decomp_len));
}

/* Test 6: streaming compression with small out buffer (NO_FLUSH then FINISH) */
void test_deflate_streaming_small_outbuf(void) {
    /* Prepare repetitive input to get decent compression */
    unsigned char input[1024];
    for (size_t i = 0; i < sizeof(input); ++i) input[i] = (unsigned char)('A' + (i % 3));

    unsigned char comp[4096];
    uLong comp_len = 0;
    int ret = compress_with_deflate_stream_two_phase(input, (uLong)sizeof(input), comp, (uLong)sizeof(comp), &comp_len);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, ret);
    TEST_ASSERT_TRUE(comp_len > 0);
    TEST_ASSERT_TRUE(comp_len < sizeof(input)); /* should compress */

    unsigned char decomp[1024];
    uLongf decomp_len = sizeof(decomp);
    int ur = uncompress(decomp, &decomp_len, comp, comp_len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ur);
    TEST_ASSERT_EQUAL_UINT(sizeof(input), decomp_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(input, decomp, decomp_len));
}

/* Test 7: After finishing, calling deflate again with Z_FINISH returns Z_STREAM_END */
void test_deflate_finish_twice_returns_stream_end(void) {
    const unsigned char input[] = "abcabcabc";
    unsigned char outbuf[256];

    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    strm.next_in = (Bytef *)input;
    strm.avail_in = (uInt)sizeof(input) - 1; /* exclude terminating null */
    strm.next_out = outbuf;
    strm.avail_out = sizeof(outbuf);

    int dret;
    do {
        dret = deflate(&strm, Z_FINISH);
        TEST_ASSERT_TRUE(dret == Z_OK || dret == Z_STREAM_END);
    } while (dret != Z_STREAM_END);

    /* Call again with Z_FINISH: should still be Z_STREAM_END */
    strm.next_out = outbuf + strm.total_out; /* set valid pointer */
    strm.avail_out = (uInt)(sizeof(outbuf) - strm.total_out);
    dret = deflate(&strm, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, dret);

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

/* Test 8: Providing input after finish should yield Z_BUF_ERROR */
void test_deflate_input_after_finish_returns_buf_error(void) {
    const unsigned char input[] = "data";
    unsigned char outbuf[128];

    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    /* Finish compression first */
    strm.next_in = (Bytef *)input;
    strm.avail_in = (uInt)sizeof(input);
    strm.next_out = outbuf;
    strm.avail_out = sizeof(outbuf);

    int dret;
    do {
        dret = deflate(&strm, Z_FINISH);
        TEST_ASSERT_TRUE(dret == Z_OK || dret == Z_STREAM_END);
    } while (dret != Z_STREAM_END);

    /* Now attempt to provide new input after finish */
    unsigned char more = 0x42;
    strm.next_in = &more;
    strm.avail_in = 1;
    strm.next_out = outbuf + strm.total_out;
    strm.avail_out = (uInt)(sizeof(outbuf) - strm.total_out);

    int dret2 = deflate(&strm, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, dret2);

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

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflate_null_stream_error);
    RUN_TEST(test_deflate_invalid_flush_value);
    RUN_TEST(test_deflate_null_next_out);
    RUN_TEST(test_deflate_zero_avail_out);
    RUN_TEST(test_deflate_roundtrip_small);
    RUN_TEST(test_deflate_streaming_small_outbuf);
    RUN_TEST(test_deflate_finish_twice_returns_stream_end);
    RUN_TEST(test_deflate_input_after_finish_returns_buf_error);
    return UNITY_END();
}