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

/* The wrapped target function is defined in deflate.c as:
 *   block_state test_deflate_slow(deflate_state *s, int flush);
 * We declare it here using the opaque internal_state pointer and int return. */
extern int test_deflate_slow(struct internal_state *s, int flush);

/* Convenience numeric values matching the block_state enum order in deflate.c */
#define BS_NEED_MORE        0
#define BS_BLOCK_DONE       1
#define BS_FINISH_STARTED   2
#define BS_FINISH_DONE      3

void setUp(void) {
    /* no-op */
}

void tearDown(void) {
    /* no-op */
}

static void init_stream(z_stream *strm, int level) {
    memset(strm, 0, sizeof(*strm));
    /* Use deflateInit2_ so we control parameters explicitly. */
    int ret = deflateInit2_(strm,
                            level,            /* level triggers slow strategy paths */
                            Z_DEFLATED,
                            MAX_WBITS,        /* 15 */
                            DEF_MEM_LEVEL,    /* 8 */
                            Z_DEFAULT_STRATEGY,
                            ZLIB_VERSION,
                            (int)sizeof(z_stream));
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

/* Test: With no input and Z_NO_FLUSH, deflate_slow should request more input (need_more). */
void test_deflate_slow_no_input_need_more(void) {
    z_stream strm;
    init_stream(&strm, 6);

    /* No input provided */
    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    /* Output buffer (won't be used in this case) */
    unsigned char outbuf[64];
    strm.next_out = outbuf;
    strm.avail_out = sizeof(outbuf);

    /* Call the wrapped target */
    int bs = test_deflate_slow((struct internal_state *)strm.state, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(BS_NEED_MORE, bs);

    deflateEnd(&strm);
}

/* Test: With no input, Z_FINISH, and ample output space, deflate_slow should finish (finish_done). */
void test_deflate_slow_finish_no_input_finish_done(void) {
    z_stream strm;
    init_stream(&strm, 6);

    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    unsigned char outbuf[1024];
    strm.next_out = outbuf;
    strm.avail_out = sizeof(outbuf);

    int bs = test_deflate_slow((struct internal_state *)strm.state, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_FINISH_DONE, bs);

    /* Some trailer/header bytes should have been produced. */
    TEST_ASSERT_TRUE(strm.total_out > 0);

    deflateEnd(&strm);
}

/* Test: Provide input with ample output buffer and Z_NO_FLUSH.
 * deflate_slow should consume input and complete a block (block_done). */
void test_deflate_slow_with_input_block_done(void) {
    z_stream strm;
    init_stream(&strm, 9); /* highest level uses slow path */

    /* Prepare repetitive input to exercise matches */
    unsigned char inbuf[1024];
    for (size_t i = 0; i < sizeof(inbuf); ++i) inbuf[i] = (unsigned char)('A' + (i % 1));

    unsigned char outbuf[16384];
    memset(outbuf, 0, sizeof(outbuf));

    strm.next_in = inbuf;
    strm.avail_in = (uInt)sizeof(inbuf);
    strm.next_out = outbuf;
    strm.avail_out = (uInt)sizeof(outbuf);

    int bs = test_deflate_slow((struct internal_state *)strm.state, Z_NO_FLUSH);

    TEST_ASSERT_EQUAL_INT(BS_BLOCK_DONE, bs);
    /* All input passed to deflate_slow should have been consumed into the window. */
    TEST_ASSERT_EQUAL_UINT(0u, strm.avail_in);
    /* Some compressed output should have been produced. */
    TEST_ASSERT_TRUE(strm.total_out > 0);

    deflateEnd(&strm);
}

/* Test: Z_FINISH with zero avail_out should report finish_started (i.e., needs more output space). */
void test_deflate_slow_finish_with_no_output_space(void) {
    z_stream strm;
    init_stream(&strm, 6);

    /* No input */
    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    unsigned char outbuf[1];
    strm.next_out = outbuf;
    strm.avail_out = 0; /* Force no output space */

    int bs = test_deflate_slow((struct internal_state *)strm.state, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_FINISH_STARTED, bs);

    deflateEnd(&strm);
}

/* Test: With some input but tiny output buffer and Z_NO_FLUSH, deflate_slow should
 * indicate need_more due to insufficient output space during a block flush. */
void test_deflate_slow_input_tiny_output_need_more(void) {
    z_stream strm;
    init_stream(&strm, 6);

    unsigned char inbuf[256];
    memset(inbuf, 'B', sizeof(inbuf));

    unsigned char outbuf[1]; /* Intentionally tiny */
    strm.next_in = inbuf;
    strm.avail_in = (uInt)sizeof(inbuf);
    strm.next_out = outbuf;
    strm.avail_out = (uInt)sizeof(outbuf);

    int bs = test_deflate_slow((struct internal_state *)strm.state, Z_NO_FLUSH);

    /* The function should early-return need_more due to lack of output space during flush. */
    TEST_ASSERT_EQUAL_INT(BS_NEED_MORE, bs);
    /* We expect at least one byte to have been emitted. */
    TEST_ASSERT_TRUE(strm.total_out <= 1);

    deflateEnd(&strm);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflate_slow_no_input_need_more);
    RUN_TEST(test_deflate_slow_finish_no_input_finish_done);
    RUN_TEST(test_deflate_slow_with_input_block_done);
    RUN_TEST(test_deflate_slow_finish_with_no_output_space);
    RUN_TEST(test_deflate_slow_input_tiny_output_need_more);
    return UNITY_END();
}