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

/* Access the internal deflate_state pointer type without redefining internals */
typedef struct internal_state deflate_state;

/* Wrapper provided by the module under test */
extern int test_deflate_huff(deflate_state *s, int flush);

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

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

/* Helper to initialize a z_stream with large output buffer */
static void init_stream_with_buffers(z_stream *strm,
                                     const unsigned char *in, size_t in_len,
                                     unsigned char *out, size_t out_len) {
    memset(strm, 0, sizeof(*strm));
    /* Initialize with default settings */
    int ret = deflateInit(strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    /* Provide input and output buffers */
    strm->next_in = (Bytef *)(uintptr_t)in;
    strm->avail_in = (uInt)in_len;

    strm->next_out = out;
    strm->avail_out = (uInt)out_len;
}

/* Convenience constants for block_state integer values */
enum {
    BS_need_more = 0,
    BS_block_done = 1,
    BS_finish_started = 2,
    BS_finish_done = 3
};

/* Test: With no input and Z_NO_FLUSH, deflate_huff should request more input (need_more). */
void test_deflate_huff_need_more_without_input(void) {
    z_stream strm;
    unsigned char out[4096];
    init_stream_with_buffers(&strm, NULL, 0, out, sizeof(out));

    /* Acquire internal state pointer */
    deflate_state *s = (deflate_state *)strm.state;

    int st = test_deflate_huff(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(BS_need_more, st);

    /* No input consumed, no output produced */
    TEST_ASSERT_EQUAL_UINT32(0u, strm.total_in);
    /* total_out may remain zero (empty block not emitted on NO_FLUSH) */
    /* Clean up */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

/* Test: With small literal input and Z_NO_FLUSH, deflate_huff processes literals and returns block_done. */
void test_deflate_huff_block_done_with_literals(void) {
    const unsigned char input[] = "Hello";
    z_stream strm;
    unsigned char out[4096];
    init_stream_with_buffers(&strm, input, sizeof(input) - 1, out, sizeof(out));

    deflate_state *s = (deflate_state *)strm.state;

    int st = test_deflate_huff(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(BS_block_done, st);

    /* All input should be consumed into the window/pending buffers */
    TEST_ASSERT_EQUAL_UINT32(sizeof(input) - 1, strm.total_in);
    /* We did not request a flush, so it's possible no bytes were emitted yet */
    /* Clean up */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

/* Test: With no input and Z_FINISH, deflate_huff should emit an empty block and return finish_done. */
void test_deflate_huff_finish_empty_input(void) {
    z_stream strm;
    unsigned char out[4096];
    init_stream_with_buffers(&strm, NULL, 0, out, sizeof(out));

    deflate_state *s = (deflate_state *)strm.state;

    int st = test_deflate_huff(s, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_finish_done, st);

    /* No input consumed */
    TEST_ASSERT_EQUAL_UINT32(0u, strm.total_in);
    /* Should have produced at least a small amount of output for the empty block */
    TEST_ASSERT_TRUE(strm.total_out > 0);

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

/* Test: With small literal input and Z_FINISH, deflate_huff should consume input, flush, and return finish_done. */
void test_deflate_huff_finish_with_literals(void) {
    const unsigned char input[] = "ABCXYZ123";
    z_stream strm;
    unsigned char out[8192];
    init_stream_with_buffers(&strm, input, sizeof(input) - 1, out, sizeof(out));

    deflate_state *s = (deflate_state *)strm.state;

    int st = test_deflate_huff(s, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_finish_done, st);

    /* All input should be consumed */
    TEST_ASSERT_EQUAL_UINT32(sizeof(input) - 1, strm.total_in);
    /* And some output should be produced due to finishing and flushing the block */
    TEST_ASSERT_TRUE(strm.total_out > 0);

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

/* Test: Two-step process — first NO_FLUSH to process literals, then FINISH to complete */
void test_deflate_huff_two_step_no_flush_then_finish(void) {
    const unsigned char input1[] = "ChunkOne-";
    const unsigned char input2[] = "ChunkTwo";
    unsigned char out[8192];

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

    /* Provide output buffer once */
    strm.next_out = out;
    strm.avail_out = (uInt)sizeof(out);

    /* First chunk with NO_FLUSH */
    strm.next_in = (Bytef *)(uintptr_t)input1;
    strm.avail_in = (uInt)(sizeof(input1) - 1);
    deflate_state *s = (deflate_state *)strm.state;
    int st1 = test_deflate_huff(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(BS_block_done, st1);
    TEST_ASSERT_EQUAL_UINT32(sizeof(input1) - 1, strm.total_in);

    /* Second chunk with FINISH */
    strm.next_in = (Bytef *)(uintptr_t)input2;
    strm.avail_in = (uInt)(sizeof(input2) - 1);
    int st2 = test_deflate_huff(s, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(BS_finish_done, st2);

    /* Both chunks should be consumed */
    TEST_ASSERT_EQUAL_UINT32((sizeof(input1) - 1) + (sizeof(input2) - 1), strm.total_in);
    /* Output should be present */
    TEST_ASSERT_TRUE(strm.total_out > 0);

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

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflate_huff_need_more_without_input);
    RUN_TEST(test_deflate_huff_block_done_with_literals);
    RUN_TEST(test_deflate_huff_finish_empty_input);
    RUN_TEST(test_deflate_huff_finish_with_literals);
    RUN_TEST(test_deflate_huff_two_step_no_flush_then_finish);
    return UNITY_END();
}