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

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

/*
 * Custom allocator to simulate allocation failures specifically for deflateCopy.
 */
static int g_fail_alloc_enabled = 0;
static size_t g_fail_alloc_size = 0; /* when enabled, fail if items*size == this */

static voidpf test_zalloc(voidpf opaque, uInt items, uInt size) {
    (void)opaque;
    size_t bytes = (size_t)items * (size_t)size;
    if (g_fail_alloc_enabled && bytes == g_fail_alloc_size) {
        return NULL;
    }
    return calloc(items, size);
}

static void test_zfree(voidpf opaque, voidpf address) {
    (void)opaque;
    free(address);
}

/* Helper: compress given input with provided stream and flush, capturing only
 * the output generated during this call sequence. Returns zlib ret code from
 * the last deflate() call. Allocates *outbuf; caller must free it. */
static int deflate_with_input(z_stream *strm, const unsigned char *in, uInt in_len,
                              int flush, unsigned char **outbuf, size_t *outlen) {
    *outbuf = NULL;
    *outlen = 0;

    /* Start with a reasonable bound; allow growth if needed. */
    uLong bound = deflateBound(strm, (uLong)in_len);
    size_t cap = (size_t)bound + 64; /* small extra cushion */
    unsigned char *buf = (unsigned char *)malloc(cap);
    if (!buf) return Z_MEM_ERROR;

    size_t produced_start = (size_t)strm->total_out;

    strm->next_in = (Bytef *)in;
    strm->avail_in = in_len;
    strm->next_out = buf;
    strm->avail_out = (uInt)cap;

    int ret;
    for (;;) {
        ret = deflate(strm, flush);
        if (ret == Z_STREAM_ERROR) {
            free(buf);
            return ret;
        }
        if (strm->avail_in == 0 && (flush != Z_FINISH || ret == Z_STREAM_END)) {
            break;
        }
        if (strm->avail_out == 0) {
            /* grow buffer */
            size_t used = (size_t)(strm->total_out - produced_start);
            size_t new_cap = cap * 2 + 1024;
            unsigned char *new_buf = (unsigned char *)realloc(buf, new_cap);
            if (!new_buf) {
                free(buf);
                return Z_MEM_ERROR;
            }
            buf = new_buf;
            cap = new_cap;
            strm->next_out = buf + used;
            /* avail_out is what's left */
            strm->avail_out = (uInt)(cap - used);
        }
    }

    size_t produced = (size_t)(strm->total_out - produced_start);
    unsigned char *final_buf = (unsigned char *)realloc(buf, produced ? produced : 1);
    if (!final_buf && produced) {
        free(buf);
        return Z_MEM_ERROR;
    }
    *outbuf = final_buf ? final_buf : buf; /* if produced==0, keep original */
    *outlen = produced;
    return ret;
}

/* Helper: generate repeated pattern data of given total size. Caller frees. */
static unsigned char *gen_repeated(const char *pat, size_t total, size_t *out_len) {
    size_t pat_len = strlen(pat);
    if (pat_len == 0) pat_len = 1;
    unsigned char *buf = (unsigned char *)malloc(total ? total : 1);
    if (!buf) return NULL;
    for (size_t i = 0; i < total; ++i) {
        buf[i] = (unsigned char)pat[i % pat_len];
    }
    if (out_len) *out_len = total;
    return buf;
}

void setUp(void) {
    /* Reset allocator failure flags before each test */
    g_fail_alloc_enabled = 0;
    g_fail_alloc_size = 0;
}

void tearDown(void) {
    /* Nothing */
}

/* Test: dest == NULL -> Z_STREAM_ERROR */
void test_deflateCopy_null_dest(void) {
    z_stream src;
    memset(&src, 0, sizeof(src));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit(&src, Z_DEFAULT_COMPRESSION));

    int ret = deflateCopy(NULL, &src);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

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

/* Test: source == NULL -> Z_STREAM_ERROR */
void test_deflateCopy_null_source(void) {
    z_stream dest;
    memset(&dest, 0, sizeof(dest));
    int ret = deflateCopy(&dest, NULL);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: uninitialized source (zeroed) -> Z_STREAM_ERROR */
void test_deflateCopy_uninitialized_source(void) {
    z_stream src, dest;
    memset(&src, 0, sizeof(src));
    memset(&dest, 0, sizeof(dest));
    int ret = deflateCopy(&dest, &src);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Test: successful copy mid-stream yields identical subsequent output. */
void test_deflateCopy_success_equal_outputs(void) {
    /* Prepare data */
    size_t len1, len2;
    unsigned char *data1 = gen_repeated("The quick brown fox jumps over the lazy dog\n", 4096, &len1);
    unsigned char *data2 = gen_repeated("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", 4096, &len2);
    TEST_ASSERT_NOT_NULL(data1);
    TEST_ASSERT_NOT_NULL(data2);

    z_stream src;
    memset(&src, 0, sizeof(src));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit(&src, 6));

    /* Feed first part and flush to boundary so pending == 0 */
    unsigned char *out1 = NULL; size_t out1_len = 0;
    int ret = deflate_with_input(&src, data1, (uInt)len1, Z_SYNC_FLUSH, &out1, &out1_len);
    TEST_ASSERT(ret == Z_OK || ret == Z_BUF_ERROR); /* Z_OK typical; Z_BUF_ERROR shouldn't occur here but allow */

    unsigned pending = 0; int bits = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&src, &pending, &bits));
    TEST_ASSERT_EQUAL_UINT(0u, pending);

    /* Copy the stream */
    z_stream dst;
    /* Ensure dst has no junk that could interfere */
    memset(&dst, 0, sizeof(dst));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateCopy(&dst, &src));

    /* Now finish both with the same tail input and compare the outputs */
    unsigned char *out2_src = NULL, *out2_dst = NULL;
    size_t out2_src_len = 0, out2_dst_len = 0;

    int ret_src = deflate_with_input(&src, data2, (uInt)len2, Z_FINISH, &out2_src, &out2_src_len);
    int ret_dst = deflate_with_input(&dst, data2, (uInt)len2, Z_FINISH, &out2_dst, &out2_dst_len);

    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, ret_src);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, ret_dst);
    TEST_ASSERT_EQUAL_size_t(out2_src_len, out2_dst_len);
    TEST_ASSERT_EQUAL_INT(0, memcmp(out2_src, out2_dst, out2_src_len));

    /* Cleanup */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&src));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&dst));

    free(data1);
    free(data2);
    free(out1);
    free(out2_src);
    free(out2_dst);
}

/* Test: simulate allocation failure on first deflate_state allocation in deflateCopy -> Z_MEM_ERROR
 * and confirm source remains usable.
 */
void test_deflateCopy_alloc_failure_on_state(void) {
    /* Initialize source with custom allocator that we can toggle to fail later */
    z_stream src;
    memset(&src, 0, sizeof(src));
    src.zalloc = test_zalloc;
    src.zfree = test_zfree;

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit(&src, Z_DEFAULT_COMPRESSION));

    /* Toggle failure for allocation size of deflate_state */
    g_fail_alloc_enabled = 1;
    g_fail_alloc_size = sizeof(struct internal_state); /* deflate_state size */

    z_stream dest;
    memset(&dest, 0, sizeof(dest));

    int ret = deflateCopy(&dest, &src);
    TEST_ASSERT_EQUAL_INT(Z_MEM_ERROR, ret);

    /* Ensure source still works after the failed copy */
    const char *msg = "hello world";
    unsigned char *out = NULL; size_t out_len = 0;
    int r2 = deflate_with_input(&src, (const unsigned char *)msg, (uInt)strlen(msg), Z_FINISH, &out, &out_len);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, r2);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&src));
    free(out);

    /* reset allocator flags */
    g_fail_alloc_enabled = 0;
    g_fail_alloc_size = 0;
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflateCopy_null_dest);
    RUN_TEST(test_deflateCopy_null_source);
    RUN_TEST(test_deflateCopy_uninitialized_source);
    RUN_TEST(test_deflateCopy_success_equal_outputs);
    RUN_TEST(test_deflateCopy_alloc_failure_on_state);
    return UNITY_END();
}