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

/* The function under test is local in deflate.c; a global wrapper is provided. */
extern unsigned test_read_buf(z_streamp strm, Bytef *buf, unsigned size);

static z_stream g_strm;

static void init_stream_with_wrap(int wrap_zlib)
{
    memset(&g_strm, 0, sizeof(g_strm));
    int windowBits = wrap_zlib ? 15 : -15; /* 15 -> zlib wrapper (wrap==1), -15 -> raw (wrap==0) */
    int ret = deflateInit2_(&g_strm,
                            Z_DEFAULT_COMPRESSION,
                            Z_DEFLATED,
                            windowBits,
                            DEF_MEM_LEVEL,
                            Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION,
                            (int)sizeof(z_stream));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

void setUp(void) {
    /* No-op; tests initialize the stream per-case. */
}

void tearDown(void) {
    if (g_strm.state != Z_NULL) {
        (void)deflateEnd(&g_strm);
        memset(&g_strm, 0, sizeof(g_strm));
    }
}

static void set_input(const unsigned char *data, unsigned len)
{
    g_strm.next_in = (z_const Bytef *)data;
    g_strm.avail_in = len;
    g_strm.total_in = 0;
}

void test_read_buf_zero_avail_returns_zero_and_no_changes(void)
{
    init_stream_with_wrap(1); /* wrap == 1 (zlib) */
    const unsigned char input[] = {0x10, 0x20, 0x30};
    set_input(input, 0); /* no available input */

    /* Pre-set fields to verify no changes occur */
    const z_const Bytef *orig_next_in = g_strm.next_in;
    g_strm.total_in = 7;
    g_strm.adler = 0xDEADBEEF;

    unsigned char out[8];
    memset(out, 0xAA, sizeof(out));

    unsigned got = test_read_buf(&g_strm, out, 5);
    TEST_ASSERT_EQUAL_UINT(0, got);
    TEST_ASSERT_EQUAL_UINT(0, g_strm.avail_in);
    TEST_ASSERT_EQUAL_PTR(orig_next_in, g_strm.next_in);
    TEST_ASSERT_EQUAL_UINT(7u, g_strm.total_in);
    TEST_ASSERT_EQUAL_UINT32(0xDEADBEEF, (uint32_t)g_strm.adler);
    TEST_ASSERT_EQUAL_UINT8(0xAA, out[0]); /* buffer untouched */
}

void test_read_buf_partial_copy_updates_counters_and_adler_wrap1(void)
{
    init_stream_with_wrap(1); /* wrap == 1 (zlib) */
    const unsigned char input[] = {'a','b','c','d','e','f','g','h','i'}; /* 9 bytes */
    set_input(input, (unsigned)sizeof(input));

    /* Set a known starting adler */
    g_strm.adler = adler32(0L, Z_NULL, 0); /* 1 */

    const z_const Bytef *orig_next_in = g_strm.next_in;
    unsigned orig_avail = g_strm.avail_in;

    unsigned char out[16];
    memset(out, 0, sizeof(out));

    unsigned req = 4;
    unsigned got = test_read_buf(&g_strm, out, req);
    TEST_ASSERT_EQUAL_UINT(req, got);

    /* Output content */
    TEST_ASSERT_EQUAL_UINT8_ARRAY(input, out, req);

    /* Counters and pointers */
    TEST_ASSERT_EQUAL_UINT(orig_avail - req, g_strm.avail_in);
    TEST_ASSERT_EQUAL_PTR(orig_next_in + req, g_strm.next_in);
    TEST_ASSERT_EQUAL_UINT(req, g_strm.total_in);

    /* Adler updated over the bytes actually read */
    uLong expected_adler = adler32(adler32(0L, Z_NULL, 0), out, req);
    TEST_ASSERT_EQUAL_UINT32(expected_adler, (uint32_t)g_strm.adler);
}

void test_read_buf_size_exceeds_avail_reads_all_and_updates_adler(void)
{
    init_stream_with_wrap(1); /* wrap == 1 (zlib) */
    const unsigned char input[] = {1,2,3,4,5,6}; /* 6 bytes */
    set_input(input, (unsigned)sizeof(input));
    g_strm.adler = adler32(0L, Z_NULL, 0);

    const z_const Bytef *orig_next_in = g_strm.next_in;

    unsigned char out[32];
    memset(out, 0, sizeof(out));

    unsigned got = test_read_buf(&g_strm, out, 20); /* request more than avail */
    TEST_ASSERT_EQUAL_UINT(sizeof(input), got);

    TEST_ASSERT_EQUAL_UINT8_ARRAY(input, out, (unsigned)sizeof(input));

    TEST_ASSERT_EQUAL_UINT(0, g_strm.avail_in);
    TEST_ASSERT_EQUAL_PTR(orig_next_in + sizeof(input), g_strm.next_in);
    TEST_ASSERT_EQUAL_UINT(sizeof(input), g_strm.total_in);

    uLong expected_adler = adler32(adler32(0L, Z_NULL, 0), out, (unsigned)sizeof(input));
    TEST_ASSERT_EQUAL_UINT32(expected_adler, (uint32_t)g_strm.adler);
}

void test_read_buf_wrap0_does_not_update_adler(void)
{
    init_stream_with_wrap(0); /* wrap == 0 (raw) */
    const unsigned char input[] = {9,8,7,6,5};
    set_input(input, (unsigned)sizeof(input));

    g_strm.adler = 0x11223344; /* sentinel */

    unsigned char out[8];
    memset(out, 0, sizeof(out));

    unsigned got = test_read_buf(&g_strm, out, (unsigned)sizeof(input));
    TEST_ASSERT_EQUAL_UINT(sizeof(input), got);

    TEST_ASSERT_EQUAL_UINT8_ARRAY(input, out, (unsigned)sizeof(input));
    TEST_ASSERT_EQUAL_UINT32(0x11223344, (uint32_t)g_strm.adler); /* unchanged */
}

void test_read_buf_multiple_calls_accumulate_total_in_and_adler(void)
{
    init_stream_with_wrap(1); /* wrap == 1 (zlib) */
    const unsigned char input[] = {10,20,30,40,50,60,70,80,90,100}; /* 10 bytes */
    set_input(input, (unsigned)sizeof(input));

    g_strm.adler = adler32(0L, Z_NULL, 0);

    unsigned char out1[4] = {0};
    unsigned char out2[4] = {0};

    unsigned got1 = test_read_buf(&g_strm, out1, 3);
    TEST_ASSERT_EQUAL_UINT(3, got1);

    uLong expected_adler = adler32(adler32(0L, Z_NULL, 0), out1, got1);
    TEST_ASSERT_EQUAL_UINT32(expected_adler, (uint32_t)g_strm.adler);
    TEST_ASSERT_EQUAL_UINT(3, g_strm.total_in);
    TEST_ASSERT_EQUAL_UINT(sizeof(input) - 3, g_strm.avail_in);

    unsigned got2 = test_read_buf(&g_strm, out2, 4);
    TEST_ASSERT_EQUAL_UINT(4, got2);

    expected_adler = adler32(expected_adler, out2, got2);
    TEST_ASSERT_EQUAL_UINT32(expected_adler, (uint32_t)g_strm.adler);
    TEST_ASSERT_EQUAL_UINT(7, g_strm.total_in);
    TEST_ASSERT_EQUAL_UINT(sizeof(input) - 7, g_strm.avail_in);

    /* Verify the data copied in both calls is exactly the first 7 bytes */
    TEST_ASSERT_EQUAL_UINT8_ARRAY(input, out1, got1);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(input + 3, out2, got2);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_read_buf_zero_avail_returns_zero_and_no_changes);
    RUN_TEST(test_read_buf_partial_copy_updates_counters_and_adler_wrap1);
    RUN_TEST(test_read_buf_size_exceeds_avail_reads_all_and_updates_adler);
    RUN_TEST(test_read_buf_wrap0_does_not_update_adler);
    RUN_TEST(test_read_buf_multiple_calls_accumulate_total_in_and_adler);
    return UNITY_END();
}