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

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

/* Wrapper provided in the source for local function testing. Avoid including internal headers. */
extern int test_deflate_stored(void *s, int flush);

static void fill_pattern(unsigned char *buf, size_t len) {
    /* Simple deterministic pattern that's not all the same byte */
    for (size_t i = 0; i < len; i++) {
        buf[i] = (unsigned char)((i * 31u + 7u) & 0xFF);
    }
}

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

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

/* Helper to initialize a deflate stream for testing. */
static void init_stream(z_stream *strm, int level) {
    memset(strm, 0, sizeof(*strm));
    /* Use zlib wrapper (wrap==1), standard window and memLevel. */
    int ret = deflateInit2(strm, level, Z_DEFLATED, MAX_WBITS, 8, Z_DEFAULT_STRATEGY);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

/* Helper to clean up a deflate stream */
static void cleanup_stream(z_stream *strm) {
    int ret = deflateEnd(strm);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
}

/* Returns the number of data bytes in out that match the input tail, by
   deducing the header length as total_out - input_len. Ensures header_len>=0. */
static void assert_output_has_input_tail(const unsigned char *out, uLong total_out,
                                         const unsigned char *in, size_t in_len) {
    TEST_ASSERT_TRUE_MESSAGE(total_out >= (uLong)in_len, "total_out is smaller than input length; no room for header");
    size_t header_len = (size_t)(total_out - (uLong)in_len);
    /* Sanity: stored block header is at least a few bytes; do a loose check */
    TEST_ASSERT_TRUE_MESSAGE(header_len >= 1, "Header length should be at least 1 byte");
    TEST_ASSERT_EQUAL_UINT8_ARRAY(in, out + header_len, in_len);
}

void test_deflate_stored_direct_copy_no_flush(void) {
    z_stream strm;
    init_stream(&strm, 0); /* level 0 to match stored behavior (though not strictly required for the direct call) */

    const size_t in_len = 1000;
    unsigned char *in = (unsigned char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in);
    fill_pattern(in, in_len);

    const size_t out_cap = in_len + 128;
    unsigned char *out = (unsigned char *)malloc(out_cap);
    TEST_ASSERT_NOT_NULL(out);
    memset(out, 0xA5, out_cap);

    strm.next_in = in;
    strm.avail_in = (uInt)in_len;
    strm.next_out = out;
    strm.avail_out = (uInt)out_cap;

    /* Call the target function */
    int bstate = test_deflate_stored(strm.state, Z_NO_FLUSH);

    /* Expect need_more (0) since we're not flushing or finishing */
    TEST_ASSERT_EQUAL_INT(0, bstate);

    /* All input should be consumed */
    TEST_ASSERT_EQUAL_UINT(0, strm.avail_in);

    /* Output should include a stored block header plus the raw input */
    TEST_ASSERT_TRUE(strm.total_out >= (uLong)in_len);
    assert_output_has_input_tail(out, strm.total_out, in, in_len);

    /* Adler32 should reflect bytes consumed */
    uLong expected = adler32(adler32(0L, Z_NULL, 0), in, (uInt)in_len);
    TEST_ASSERT_EQUAL_UINT(expected, strm.adler);

    free(in);
    free(out);
    cleanup_stream(&strm);
}

void test_deflate_stored_finish_done_single_block(void) {
    z_stream strm;
    init_stream(&strm, 0);

    const size_t in_len = 50;
    unsigned char *in = (unsigned char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in);
    fill_pattern(in, in_len);

    const size_t out_cap = in_len + 128;
    unsigned char *out = (unsigned char *)malloc(out_cap);
    TEST_ASSERT_NOT_NULL(out);
    memset(out, 0, out_cap);

    strm.next_in = in;
    strm.avail_in = (uInt)in_len;
    strm.next_out = out;
    strm.avail_out = (uInt)out_cap;

    int bstate = test_deflate_stored(strm.state, Z_FINISH);

    /* Expect finish_done (3) since we provided enough output for the last block */
    TEST_ASSERT_EQUAL_INT(3, bstate);
    TEST_ASSERT_EQUAL_UINT(0, strm.avail_in);

    /* Output should contain header + input */
    TEST_ASSERT_TRUE(strm.total_out >= (uLong)in_len);
    assert_output_has_input_tail(out, strm.total_out, in, in_len);

    /* Adler32 reflects all input consumed */
    uLong expected = adler32(adler32(0L, Z_NULL, 0), in, (uInt)in_len);
    TEST_ASSERT_EQUAL_UINT(expected, strm.adler);

    free(in);
    free(out);
    cleanup_stream(&strm);
}

void test_deflate_stored_no_output_space_large_input(void) {
    z_stream strm;
    init_stream(&strm, 0);

    const size_t in_len = 70000; /* Larger than default window_size (64K) to exercise window filling */
    unsigned char *in = (unsigned char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in);
    fill_pattern(in, in_len);

    unsigned char out_dummy;

    strm.next_in = in;
    strm.avail_in = (uInt)in_len;
    strm.next_out = &out_dummy;
    strm.avail_out = 0; /* No room for any header, forces pending path */

    uLong total_in_before = strm.total_in;
    uLong total_out_before = strm.total_out;

    int bstate = test_deflate_stored(strm.state, Z_NO_FLUSH);

    /* Expect need_more (0) */
    TEST_ASSERT_EQUAL_INT(0, bstate);

    /* Some input should have been consumed into the window */
    TEST_ASSERT_TRUE(strm.total_in > total_in_before);

    /* No output could be written */
    TEST_ASSERT_EQUAL_UINT(total_out_before, strm.total_out);

    /* Adler32 should match the prefix that was read */
    uLong consumed = strm.total_in - total_in_before;
    uLong expected = adler32(adler32(0L, Z_NULL, 0), in, (uInt)consumed);
    TEST_ASSERT_EQUAL_UINT(expected, strm.adler);

    free(in);
    cleanup_stream(&strm);
}

void test_deflate_stored_sync_flush_block_done(void) {
    z_stream strm;
    init_stream(&strm, 0);

    const size_t in_len = 10;
    unsigned char *in = (unsigned char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in);
    fill_pattern(in, in_len);

    const size_t out_cap = in_len + 64;
    unsigned char *out = (unsigned char *)malloc(out_cap);
    TEST_ASSERT_NOT_NULL(out);
    memset(out, 0, out_cap);

    strm.next_in = in;
    strm.avail_in = (uInt)in_len;
    strm.next_out = out;
    strm.avail_out = (uInt)out_cap;

    int bstate = test_deflate_stored(strm.state, Z_SYNC_FLUSH);

    /* Expect block_done (1) when flushing and all input consumed */
    TEST_ASSERT_EQUAL_INT(1, bstate);
    TEST_ASSERT_EQUAL_UINT(0, strm.avail_in);

    /* Output should include header + input */
    TEST_ASSERT_TRUE(strm.total_out >= (uLong)in_len);
    assert_output_has_input_tail(out, strm.total_out, in, in_len);

    /* Adler32 reflects all input consumed */
    uLong expected = adler32(adler32(0L, Z_NULL, 0), in, (uInt)in_len);
    TEST_ASSERT_EQUAL_UINT(expected, strm.adler);

    free(in);
    free(out);
    cleanup_stream(&strm);
}

void test_deflate_stored_too_little_out_for_header_small_input(void) {
    z_stream strm;
    init_stream(&strm, 0);

    const size_t in_len = 10;
    unsigned char *in = (unsigned char *)malloc(in_len);
    TEST_ASSERT_NOT_NULL(in);
    fill_pattern(in, in_len);

    unsigned char out4[4];

    strm.next_in = in;
    strm.avail_in = (uInt)in_len;
    strm.next_out = out4;
    strm.avail_out = 4; /* Likely less than header bytes, prevents direct output */

    uLong total_in_before = strm.total_in;
    uLong total_out_before = strm.total_out;

    int bstate = test_deflate_stored(strm.state, Z_NO_FLUSH);

    /* Expect need_more (0) */
    TEST_ASSERT_EQUAL_INT(0, bstate);

    /* All input should have been read into the window since space is abundant there */
    TEST_ASSERT_EQUAL_UINT((uLong)(total_in_before + in_len), strm.total_in);

    /* No output written due to insufficient space for header */
    TEST_ASSERT_EQUAL_UINT(total_out_before, strm.total_out);

    /* Adler32 should reflect the full input consumed */
    uLong expected = adler32(adler32(0L, Z_NULL, 0), in, (uInt)in_len);
    TEST_ASSERT_EQUAL_UINT(expected, strm.adler);

    free(in);
    cleanup_stream(&strm);
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflate_stored_direct_copy_no_flush);
    RUN_TEST(test_deflate_stored_finish_done_single_block);
    RUN_TEST(test_deflate_stored_no_output_space_large_input);
    RUN_TEST(test_deflate_stored_sync_flush_block_done);
    RUN_TEST(test_deflate_stored_too_little_out_for_header_small_input);
    return UNITY_END();
}