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

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

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

/* Helpers */

static void fill_pattern(unsigned char *buf, size_t len, unsigned seed) {
    /* Simple deterministic pattern for verification */
    for (size_t i = 0; i < len; i++) {
        buf[i] = (unsigned char)((i * 1315423911u + seed) & 0xFF);
    }
}

/* Tests */

void test_deflateGetDictionary_null_stream_returns_error(void) {
    uInt len = 1234;
    int ret = deflateGetDictionary(NULL, NULL, &len);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

void test_deflateGetDictionary_uninitialized_stream_returns_error(void) {
    z_stream strm;
    memset(&strm, 0, sizeof(strm)); /* not initialized */
    uInt len = 0;
    int ret = deflateGetDictionary(&strm, NULL, &len);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

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

    uInt len = 0xDEAD; /* ensure it gets set to 0 */
    int ret = deflateGetDictionary(&strm, NULL, &len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_EQUAL_UINT(0u, len);

    /* Also verify call with non-NULL dictionary pointer does not crash and still reports zero length */
    unsigned char dummy[8] = {0};
    len = 0xBEEF;
    ret = deflateGetDictionary(&strm, dummy, &len);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    TEST_ASSERT_EQUAL_UINT(0u, len);

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

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

    const uInt dict_len = 1024;
    unsigned char *dict = (unsigned char *)malloc(dict_len);
    TEST_ASSERT_NOT_NULL(dict);
    fill_pattern(dict, dict_len, 7);

    /* Set dictionary */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateSetDictionary(&strm, dict, dict_len));

    /* Two-phase retrieval: first get length */
    uInt got_len = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, NULL, &got_len));
    TEST_ASSERT_EQUAL_UINT(dict_len, got_len);

    /* Retrieve the content */
    unsigned char *out = (unsigned char *)malloc(got_len ? got_len : 1);
    TEST_ASSERT_NOT_NULL(out);
    uInt got_len2 = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, out, &got_len2));
    TEST_ASSERT_EQUAL_UINT(got_len, got_len2);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(dict, out, got_len);

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

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

    /* Use a large dictionary greater than the default window (32 KiB) */
    const uInt large_len = 50000u;
    unsigned char *dict = (unsigned char *)malloc(large_len);
    TEST_ASSERT_NOT_NULL(dict);
    fill_pattern(dict, large_len, 123);

    /* Set dictionary */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateSetDictionary(&strm, dict, large_len));

    /* Get length first */
    uInt got_len = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, NULL, &got_len));

    /* Expect truncation to window size. The default windowBits is 15 (32 KiB). */
    TEST_ASSERT_TRUE_MESSAGE(got_len <= large_len, "Returned length must be <= provided dictionary length");
    TEST_ASSERT_TRUE_MESSAGE(got_len > 0, "Returned length should be > 0 for non-empty dictionary");

    /* Retrieve content */
    unsigned char *out = (unsigned char *)malloc(got_len);
    TEST_ASSERT_NOT_NULL(out);
    uInt got_len2 = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, out, &got_len2));
    TEST_ASSERT_EQUAL_UINT(got_len, got_len2);

    /* The expected data is the last got_len bytes of the original dict */
    const unsigned char *expected_tail = dict + (large_len - got_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected_tail, out, got_len);

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

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

    const uInt dict_len = 7777;
    unsigned char *dict = (unsigned char *)malloc(dict_len);
    TEST_ASSERT_NOT_NULL(dict);
    fill_pattern(dict, dict_len, 55);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateSetDictionary(&strm, dict, dict_len));

    uInt reported_len = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, NULL, &reported_len));
    TEST_ASSERT_TRUE(reported_len > 0);

    /* Retrieve again to verify content matches the expected tail */
    unsigned char *out = (unsigned char *)malloc(reported_len);
    TEST_ASSERT_NOT_NULL(out);
    uInt reported_len2 = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, out, &reported_len2));
    TEST_ASSERT_EQUAL_UINT(reported_len, reported_len2);

    const unsigned char *expected_tail = dict + (dict_len - reported_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected_tail, out, reported_len);

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

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

    const uInt dict_len = 4096;
    unsigned char *dict = (unsigned char *)malloc(dict_len);
    TEST_ASSERT_NOT_NULL(dict);
    fill_pattern(dict, dict_len, 201);

    TEST_ASSERT_EQUAL_INT(Z_OK, deflateSetDictionary(&strm, dict, dict_len));

    /* First query length so we know how many bytes to expect */
    uInt expected_len = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, NULL, &expected_len));
    TEST_ASSERT_TRUE(expected_len > 0);
    TEST_ASSERT_TRUE(expected_len <= dict_len);

    unsigned char *out = (unsigned char *)malloc(expected_len);
    TEST_ASSERT_NOT_NULL(out);

    /* Now call with dictLength == NULL; expect it still copies the bytes */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateGetDictionary(&strm, out, NULL));

    const unsigned char *expected_tail = dict + (dict_len - expected_len);
    TEST_ASSERT_EQUAL_UINT8_ARRAY(expected_tail, out, expected_len);

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

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

    /* End the stream */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflateEnd(&strm));

    /* Now calling deflateGetDictionary should fail */
    uInt len = 0;
    int ret = deflateGetDictionary(&strm, NULL, &len);
    TEST_ASSERT_EQUAL_INT(Z_STREAM_ERROR, ret);
}

/* Unity main */
int main(void) {
    UNITY_BEGIN();

    RUN_TEST(test_deflateGetDictionary_null_stream_returns_error);
    RUN_TEST(test_deflateGetDictionary_uninitialized_stream_returns_error);
    RUN_TEST(test_deflateGetDictionary_no_dictionary_set_returns_zero_length);
    RUN_TEST(test_deflateGetDictionary_round_trip_small_dictionary);
    RUN_TEST(test_deflateGetDictionary_large_dictionary_truncated_to_window);
    RUN_TEST(test_deflateGetDictionary_null_buffer_only_length);
    RUN_TEST(test_deflateGetDictionary_null_length_pointer_copies_data);
    RUN_TEST(test_deflateGetDictionary_after_deflateEnd_returns_error);

    return UNITY_END();
}