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

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

/* Prototype for the auto-generated wrapper in gzread.c */
int test_gz_decomp(gz_statep state);

/* Helpers */

static int compress_to_gzip(const unsigned char *in, size_t inlen,
                            unsigned char **out, size_t *outlen) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit2(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED,
                           15 + 16, 8, Z_DEFAULT_STRATEGY); /* gzip wrapper */
    if (ret != Z_OK) return ret;

    uLong bound = deflateBound(&strm, (uLong)inlen);
    unsigned char *buf = (unsigned char *)malloc(bound);
    if (!buf) {
        deflateEnd(&strm);
        return Z_MEM_ERROR;
    }

    strm.next_in = (Bytef *)in;
    strm.avail_in = (uInt)inlen;
    strm.next_out = buf;
    strm.avail_out = (uInt)bound;

    ret = deflate(&strm, Z_FINISH);
    if (ret != Z_STREAM_END) {
        free(buf);
        deflateEnd(&strm);
        return ret == Z_OK ? Z_BUF_ERROR : ret;
    }

    *outlen = (size_t)((Bytef *)strm.next_out - buf);
    *out = buf;

    deflateEnd(&strm);
    return Z_OK;
}

static gz_statep make_state_with_input(const unsigned char *comp, unsigned comp_len) {
    struct gz_state *state = (struct gz_state *)calloc(1, sizeof(struct gz_state));
    if (!state) return NULL;

    state->how = GZIP;      /* we explicitly want to run through gzip inflate */
    state->direct = 0;
    state->err = Z_OK;
    state->msg = NULL;
    state->eof = 0;
    state->x.have = 0;
    state->x.next = NULL;

    z_streamp strm = &state->strm;
    memset(strm, 0, sizeof(*strm));
    strm->zalloc = Z_NULL;
    strm->zfree = Z_NULL;
    strm->opaque = Z_NULL;
    if (inflateInit2(strm, 15 + 16) != Z_OK) {
        free(state);
        return NULL;
    }

    strm->next_in = (Bytef *)comp;
    strm->avail_in = comp_len;

    return state;
}

static void destroy_state(gz_statep state) {
    if (!state) return;
    inflateEnd(&(state->strm));
    free(state);
}

/* Unity fixtures */

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

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

/* Tests */

static void test_gz_decomp_success_full_stream_end_sets_LOOK(void) {
    const char *msg = "Hello, gz_decomp via zlib! 0123456789.\n";
    const size_t msg_len = strlen(msg);

    unsigned char *gzbuf = NULL;
    size_t gzlen = 0;
    int rc = compress_to_gzip((const unsigned char *)msg, msg_len, &gzbuf, &gzlen);
    TEST_ASSERT_EQUAL_INT(Z_OK, rc);
    TEST_ASSERT_NOT_NULL(gzbuf);
    TEST_ASSERT_TRUE(gzlen > 0);

    gz_statep state = make_state_with_input(gzbuf, (unsigned)gzlen);
    TEST_ASSERT_NOT_NULL(state);

    /* Provide ample output space */
    unsigned char outbuf[1024];
    memset(outbuf, 0xA5, sizeof(outbuf));
    state->strm.next_out = outbuf;
    state->strm.avail_out = (uInt)sizeof(outbuf);

    rc = test_gz_decomp(state);
    TEST_ASSERT_EQUAL_INT(0, rc);

    /* Expect full message produced and end-of-stream moves to LOOK */
    TEST_ASSERT_EQUAL_UINT(msg_len, state->x.have);
    TEST_ASSERT_EQUAL_INT(LOOK, state->how);
    TEST_ASSERT_NOT_NULL(state->x.next);
    TEST_ASSERT_EQUAL_UINT8_ARRAY((const uint8_t *)msg, state->x.next, state->x.have);

    destroy_state(state);
    free(gzbuf);
}

static void test_gz_decomp_partial_output_no_end(void) {
    /* Make a message larger than the output buffer we will provide */
    char msg[256];
    for (int i = 0; i < (int)sizeof(msg); i++) {
        msg[i] = (char)('A' + (i % 26));
    }
    const size_t msg_len = sizeof(msg);

    unsigned char *gzbuf = NULL;
    size_t gzlen = 0;
    int rc = compress_to_gzip((const unsigned char *)msg, msg_len, &gzbuf, &gzlen);
    TEST_ASSERT_EQUAL_INT(Z_OK, rc);
    TEST_ASSERT_TRUE(gzlen > 0);

    gz_statep state = make_state_with_input(gzbuf, (unsigned)gzlen);
    TEST_ASSERT_NOT_NULL(state);

    /* Provide too small output buffer to force partial decompression */
    enum { OUT_LIMIT = 37 };
    unsigned char outbuf[OUT_LIMIT];
    memset(outbuf, 0xCD, sizeof(outbuf));
    state->strm.next_out = outbuf;
    state->strm.avail_out = (uInt)sizeof(outbuf);

    rc = test_gz_decomp(state);
    TEST_ASSERT_EQUAL_INT(0, rc);

    /* Expect exactly OUT_LIMIT bytes produced and still in GZIP mode (not at end) */
    TEST_ASSERT_EQUAL_UINT(OUT_LIMIT, state->x.have);
    TEST_ASSERT_EQUAL_INT(GZIP, state->how);
    TEST_ASSERT_NOT_NULL(state->x.next);
    TEST_ASSERT_EQUAL_UINT8_ARRAY((const uint8_t *)msg, state->x.next, state->x.have);

    destroy_state(state);
    free(gzbuf);
}

static void test_gz_decomp_unexpected_eof_sets_buf_error(void) {
    /* Create a state ready for GZIP, but with no input and eof signaled. */
    gz_statep state = make_state_with_input(NULL, 0);
    TEST_ASSERT_NOT_NULL(state);

    state->eof = 1; /* simulate EOF so gz_avail won't attempt to read */

    unsigned char outbuf[64];
    memset(outbuf, 0xEE, sizeof(outbuf));
    state->strm.next_out = outbuf;
    state->strm.avail_out = (uInt)sizeof(outbuf);

    int rc = test_gz_decomp(state);
    TEST_ASSERT_EQUAL_INT(0, rc);                 /* function returns 0 on this path */
    TEST_ASSERT_EQUAL_INT(Z_BUF_ERROR, state->err);
    TEST_ASSERT_EQUAL_UINT(0u, state->x.have);    /* no output produced */

    destroy_state(state);
}

static void test_gz_decomp_invalid_header_returns_data_error(void) {
    /* Supply invalid gzip header bytes */
    unsigned char badhdr[2] = { 0x00, 0xFF };

    gz_statep state = make_state_with_input(badhdr, sizeof(badhdr));
    TEST_ASSERT_NOT_NULL(state);

    unsigned char outbuf[32];
    state->strm.next_out = outbuf;
    state->strm.avail_out = (uInt)sizeof(outbuf);

    int rc = test_gz_decomp(state);
    TEST_ASSERT_EQUAL_INT(-1, rc);
    TEST_ASSERT_EQUAL_INT(Z_DATA_ERROR, state->err);

    destroy_state(state);
}

/* Unity main */

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_gz_decomp_success_full_stream_end_sets_LOOK);
    RUN_TEST(test_gz_decomp_partial_output_no_end);
    RUN_TEST(test_gz_decomp_unexpected_eof_sets_buf_error);
    RUN_TEST(test_gz_decomp_invalid_header_returns_data_error);
    return UNITY_END();
}