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

/* Wrapper provided in deflate.c for testing the local function */
extern int test_deflate_fast(void *s, int flush);

/* block_state numeric values as observed in deflate.c:
   need_more = 0, block_done = 1, finish_started = 2, finish_done = 3 */
enum {
    BS_need_more = 0,
    BS_block_done = 1,
    BS_finish_started = 2,
    BS_finish_done = 3
};

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

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

/* Helper: initialize a deflate stream targeting the "fast" strategy path */
static void init_def_stream(z_stream *strm, unsigned char *out_buf, size_t out_size) {
    memset(strm, 0, sizeof(*strm));
    int ret = deflateInit2(strm,
                           2,                /* level: 1..3 go through deflate_fast */
                           Z_DEFLATED,
                           MAX_WBITS,
                           DEF_MEM_LEVEL,
                           Z_DEFAULT_STRATEGY);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    strm->next_out = out_buf;
    strm->avail_out = (uInt)out_size;
}

/* Test 1: With no input and Z_NO_FLUSH, deflate_fast should request more input (need_more) */
void test_deflate_fast_need_more_when_no_input(void) {
    z_stream strm;
    unsigned char out_buf[1024];
    init_def_stream(&strm, out_buf, sizeof(out_buf));

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

    int bs = test_deflate_fast((void *)strm.state, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(BS_need_more, bs);

    TEST_ASSERT_EQUAL_UINT(0, strm.total_in);
    TEST_ASSERT_EQUAL_UINT(0, strm.avail_in);

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

/* Test 2: With sufficient input and Z_NO_FLUSH, deflate_fast processes and returns block_done */
void test_deflate_fast_block_done_with_sufficient_input(void) {
    z_stream strm;
    /* Large output buffer to avoid space limitations */
    size_t out_size = 1 << 20;
    unsigned char *out_buf = (unsigned char *)malloc(out_size);
    TEST_ASSERT_NOT_NULL(out_buf);
    init_def_stream(&strm, out_buf, out_size);

    /* Provide input > MIN_LOOKAHEAD (typically 262) to avoid early need_more */
    size_t in_len = 1024; /* safely above MIN_LOOKAHEAD */
    unsigned char *in_buf = (unsigned char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in_buf);
    for (size_t i = 0; i < in_len; i++) {
        in_buf[i] = (unsigned char)(i & 0xFF);
    }
    strm.next_in = in_buf;
    strm.avail_in = (uInt)in_len;

    int bs = test_deflate_fast((void *)strm.state, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(BS_block_done, bs);

    /* All input should be consumed by fill_window/processing */
    TEST_ASSERT_EQUAL_UINT(0, strm.avail_in);
    TEST_ASSERT_EQUAL_UINT(in_len, strm.total_in);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
    free(out_buf);
    free(in_buf);
}

/* Test 3: With Z_FINISH and sufficient output space, should return finish_done */
void test_deflate_fast_finish_done_with_enough_output(void) {
    z_stream strm;
    size_t out_size = 1 << 16; /* large enough for headers and block */
    unsigned char *out_buf = (unsigned char *)malloc(out_size);
    TEST_ASSERT_NOT_NULL(out_buf);
    init_def_stream(&strm, out_buf, out_size);

    /* Provide some small input (can also be zero) */
    const char *msg = "Finish this block quickly!";
    size_t in_len = strlen(msg);
    strm.next_in = (const Bytef *)msg;
    strm.avail_in = (uInt)in_len;

    int bs = test_deflate_fast((void *)strm.state, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_finish_done, bs);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
    free(out_buf);
}

/* Test 4: With Z_FINISH but zero output space, should return finish_started */
void test_deflate_fast_finish_started_with_no_output_space(void) {
    z_stream strm;
    unsigned char dummy; /* not used since avail_out=0 */
    init_def_stream(&strm, &dummy, 0);

    /* No input required to test empty final block behavior */
    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    int bs = test_deflate_fast((void *)strm.state, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_finish_started, bs);

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

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflate_fast_need_more_when_no_input);
    RUN_TEST(test_deflate_fast_block_done_with_sufficient_input);
    RUN_TEST(test_deflate_fast_finish_done_with_enough_output);
    RUN_TEST(test_deflate_fast_finish_started_with_no_output_space);
    return UNITY_END();
}