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

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

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

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

/* Helper: round-trip inflate for zlib-wrapped data; set dictionary if needed */
static void roundtrip_inflate_with_dict(const Bytef *comp, uLong comp_len,
                                        const Bytef *dict, uInt dict_len,
                                        const Bytef *expected, uLong expected_len) {
    z_stream istrm;
    memset(&istrm, 0, sizeof(istrm));
    int ret = inflateInit(&istrm);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    Bytef *out = (Bytef *)malloc(expected_len);
    TEST_ASSERT_NOT_NULL(out);

    istrm.next_in = (Bytef *)comp;
    istrm.avail_in = (uInt)comp_len;
    istrm.next_out = out;
    istrm.avail_out = (uInt)expected_len;

    for (;;) {
        ret = inflate(&istrm, Z_NO_FLUSH);
        if (ret == Z_NEED_DICT) {
            /* Provide the same dictionary used during deflate */
            int sret = inflateSetDictionary(&istrm, dict, dict_len);
            TEST_ASSERT_EQUAL_INT(Z_OK, sret);
            continue;
        }
        if (ret == Z_STREAM_END) break;
        TEST_ASSERT_EQUAL_INT(Z_OK, ret);
        if (istrm.avail_out == 0 && ret != Z_STREAM_END) {
            /* Should not run out of output space for expected_len */
            TEST_FAIL_MESSAGE("inflate ran out of output space");
        }
    }

    TEST_ASSERT_EQUAL_UINT64(expected_len, istrm.total_out);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected, out, (size_t)expected_len);

    inflateEnd(&istrm);
    free(out);
}

void test_deflateSetDictionary_null_parameters(void) {
    /* Case A: uninitialized stream (deflateStateCheck should fail) */
    z_stream bad = {0};
    const Bytef dictA[] = "abc";
    int ret = deflateSetDictionary(&bad, dictA, (uInt)sizeof(dictA)-1);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

    /* Case B: initialized stream but NULL dictionary */
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit(&strm, Z_DEFAULT_COMPRESSION));
    ret = deflateSetDictionary(&strm, NULL, 0);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
}

void test_deflateSetDictionary_rejects_gzip_wrap(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    /* Gzip wrapper: windowBits > 15 (e.g., 15 + 16) */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit2(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED,
                                             MAX_WBITS + 16, DEF_MEM_LEVEL, Z_DEFAULT_STRATEGY));

    const Bytef dict[] = "hello world";
    int ret = deflateSetDictionary(&strm, dict, (uInt)sizeof(dict)-1);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

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

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

    /* Cause the zlib header to be written by calling deflate with no input */
    unsigned char outbuf[64];
    memset(outbuf, 0, sizeof(outbuf));
    strm.next_out = outbuf;
    strm.avail_out = (uInt)sizeof(outbuf);
    strm.next_in = Z_NULL;
    strm.avail_in = 0;

    int dret = deflate(&strm, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(Z_OK, dret);

    const Bytef dict[] = "dict";
    int ret = deflateSetDictionary(&strm, dict, (uInt)sizeof(dict)-1);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);

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

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

    const Bytef dict[] = "hello world, this is my preset dictionary";
    const uInt dict_len = (uInt)(sizeof(dict) - 1);

    int ret = deflateSetDictionary(&strm, dict, dict_len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    /* Verify adler was updated to dictionary checksum (starting from initial adler) */
    uLong initial = adler32(0L, Z_NULL, 0);
    uLong expected_adler = adler32(initial, dict, dict_len);
    TEST_ASSERT_EQUAL_UINT32(expected_adler, strm.adler);

    /* Compress some input that will benefit from the dictionary */
    const Bytef input[] = "... world, this is my preset dictionary indeed!";
    uLong input_len = (uLong)(sizeof(input) - 1);

    uLong bound = deflateBound(&strm, input_len);
    Bytef *out = (Bytef *)malloc(bound + 32);
    TEST_ASSERT_NOT_NULL(out);

    strm.next_in = (Bytef *)input;
    strm.avail_in = (uInt)input_len;
    strm.next_out = out;
    strm.avail_out = (uInt)(bound + 32);

    int dstatus;
    do {
        dstatus = deflate(&strm, Z_FINISH);
    } while (dstatus == Z_OK);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, dstatus);

    uLong comp_len = strm.total_out;

    /* Check zlib header: preset dictionary flag set and checksum matches */
    TEST_ASSERT_TRUE_MESSAGE(comp_len >= 6, "Compressed output too small for zlib header with dictionary");
    uint8_t cmf = out[0];
    uint8_t flg = out[1];
    (void)cmf; /* cmf not strictly validated here */
    TEST_ASSERT_TRUE_MESSAGE((flg & 0x20) != 0, "Zlib FLG preset-dictionary bit not set");
    uLong header_adler = ((uLong)out[2] << 24) | ((uLong)out[3] << 16) | ((uLong)out[4] << 8) | (uLong)out[5];
    TEST_ASSERT_EQUAL_UINT32(expected_adler, header_adler);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));

    /* Round-trip inflate with dictionary */
    roundtrip_inflate_with_dict(out, comp_len, dict, dict_len, input, input_len);

    free(out);
}

void test_deflateSetDictionary_large_dictionary_adler_full(void) {
    /* Create a dictionary larger than 32K */
    const size_t big_len = 40000;
    Bytef *big_dict = (Bytef *)malloc(big_len);
    TEST_ASSERT_NOT_NULL(big_dict);
    for (size_t i = 0; i < big_len; ++i) big_dict[i] = (Bytef)(i & 0xFF);

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

    int ret = deflateSetDictionary(&strm, big_dict, (uInt)big_len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    uLong initial = adler32(0L, Z_NULL, 0);
    uLong expected_adler = adler32(initial, big_dict, (uInt)big_len);
    TEST_ASSERT_EQUAL_UINT32(expected_adler, strm.adler);

    /* Optionally compress a tiny payload to force header emission and verify header checksum */
    const Bytef input[] = "tiny";
    uLong input_len = (uLong)(sizeof(input) - 1);
    uLong bound = deflateBound(&strm, input_len);
    Bytef *out = (Bytef *)malloc(bound + 32);
    TEST_ASSERT_NOT_NULL(out);

    strm.next_in = (Bytef *)input;
    strm.avail_in = (uInt)input_len;
    strm.next_out = out;
    strm.avail_out = (uInt)(bound + 32);

    int dstatus;
    do {
        dstatus = deflate(&strm, Z_FINISH);
    } while (dstatus == Z_OK);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_END, dstatus);

    TEST_ASSERT_TRUE_MESSAGE(strm.total_out >= 6, "Compressed output too small for zlib header with dictionary");
    uLong header_adler = ((uLong)out[2] << 24) | ((uLong)out[3] << 16) | ((uLong)out[4] << 8) | (uLong)out[5];
    TEST_ASSERT_EQUAL_UINT32(expected_adler, header_adler);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));
    free(out);
    free(big_dict);
}

void test_deflateSetDictionary_raw_wrap_ok(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm));
    /* Raw deflate: negative windowBits => wrap == 0 */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateInit2(&strm, Z_DEFAULT_COMPRESSION, Z_DEFLATED,
                                             -MAX_WBITS, DEF_MEM_LEVEL, Z_DEFAULT_STRATEGY));

    const Bytef dict[] = "raw-mode-dictionary";
    uInt dict_len = (uInt)(sizeof(dict) - 1);

    /* Initial adler for raw is also set to adler32(0, NULL, 0) by deflateResetKeep */
    uLong initial_adler = strm.adler; /* should be 1 */

    int ret = deflateSetDictionary(&strm, dict, dict_len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);

    /* In raw mode, deflateSetDictionary should not modify strm.adler */
    TEST_ASSERT_EQUAL_UINT32(initial_adler, strm.adler);

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

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_deflateSetDictionary_null_parameters);
    RUN_TEST(test_deflateSetDictionary_rejects_gzip_wrap);
    RUN_TEST(test_deflateSetDictionary_rejects_after_deflate_started);
    RUN_TEST(test_deflateSetDictionary_sets_adler_and_header_and_roundtrip);
    RUN_TEST(test_deflateSetDictionary_large_dictionary_adler_full);
    RUN_TEST(test_deflateSetDictionary_raw_wrap_ok);
    return UNITY_END();
}