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

/* The wrapper provided in the module; block_state is internal, so declare as int */
extern int test_deflate_rle(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 and return its internal deflate_state* */
static deflate_state* init_deflate_state(z_stream *strm, int strategy, unsigned char *outbuf, unsigned outbuf_size) {
    memset(strm, 0, sizeof(*strm));
    int ret = deflateInit2(strm,
                           Z_DEFAULT_COMPRESSION,
                           Z_DEFLATED,
                           MAX_WBITS,
                           DEF_MEM_LEVEL,
                           strategy);
    TEST_ASSERT_EQUAL_INT(Z_OK, ret);
    /* Provide an output buffer so internal flushes can write */
    strm->next_out = outbuf;
    strm->avail_out = outbuf_size;
    /* No input provided here; tests directly manipulate the window/lookahead */
    strm->next_in = NULL;
    strm->avail_in = 0;

    deflate_state *s = (deflate_state *)strm->state;
    TEST_ASSERT_NOT_NULL(s);
    TEST_ASSERT_NOT_NULL(s->window);
    TEST_ASSERT_TRUE(s->window_size >= (ulg)(2U * s->w_size));
#ifndef LIT_MEM
    TEST_ASSERT_NOT_NULL(s->sym_buf);
#else
    TEST_ASSERT_NOT_NULL(s->d_buf);
    TEST_ASSERT_NOT_NULL(s->l_buf);
#endif
    TEST_ASSERT_NOT_NULL(s->pending_buf);
    return s;
}

/* Test: With a long run of the previous byte (distance 1), deflate_rle should
 * produce a single match of length MAX_MATCH, advancing strstart by MAX_MATCH
 * and reducing lookahead accordingly, then return need_more (0).
 */
void test_deflate_rle_consumes_long_run_and_returns_need_more(void) {
    z_stream strm;
    unsigned char out[32768];
    deflate_state *s = init_deflate_state(&strm, Z_RLE, out, sizeof(out));

    /* Prepare a long run of identical bytes 'A' so prev equals many next bytes */
    const unsigned initial_strstart = 1;
    const unsigned initial_lookahead = MAX_MATCH + 10; /* > MAX_MATCH */

    /* Fill window with 'A' from 0 up to initial_strstart + initial_lookahead */
    for (unsigned i = 0; i <= initial_strstart + initial_lookahead; i++) {
        s->window[i] = (Byte)'A';
    }
    s->strstart = initial_strstart;
    s->block_start = 0L;
    s->lookahead = initial_lookahead;
    s->match_length = 0;
    s->sym_next = 0;

    int r = test_deflate_rle(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(0, r); /* need_more */

    TEST_ASSERT_EQUAL_UINT(initial_strstart + MAX_MATCH, s->strstart);
    TEST_ASSERT_EQUAL_UINT(initial_lookahead - MAX_MATCH, s->lookahead);
    TEST_ASSERT_EQUAL_UINT(0, s->match_length);
    /* Some symbols should have been tallied (match encoded) */
    TEST_ASSERT_TRUE(s->sym_next > 0);

    deflateEnd(&strm);
}

/* Test: Non-repeating data should emit literals until lookahead == MAX_MATCH,
 * then return need_more. strstart should advance by (initial_lookahead - MAX_MATCH).
 */
void test_deflate_rle_emits_literals_until_need_more(void) {
    z_stream strm;
    unsigned char out[32768];
    deflate_state *s = init_deflate_state(&strm, Z_RLE, out, sizeof(out));

    const unsigned initial_strstart = 1;
    const unsigned initial_lookahead = MAX_MATCH + 5; /* 5 literals to process */
    /* Create strictly increasing bytes to avoid any runs */
    for (unsigned i = 0; i <= initial_strstart + initial_lookahead; i++) {
        s->window[i] = (Byte)(i & 0xFF);
    }
    s->strstart = initial_strstart;
    s->block_start = 0L;
    s->lookahead = initial_lookahead;
    s->match_length = 0;
    s->sym_next = 0;

    int r = test_deflate_rle(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(0, r); /* need_more */

    TEST_ASSERT_EQUAL_UINT(initial_strstart + (initial_lookahead - MAX_MATCH), s->strstart);
    TEST_ASSERT_EQUAL_UINT(MAX_MATCH, s->lookahead);
    TEST_ASSERT_TRUE(s->sym_next > 0); /* literals tallied */

    deflateEnd(&strm);
}

/* Test: A short run (< MIN_MATCH), e.g., two identical bytes after prev,
 * should not be encoded as a match; it should emit a literal.
 * With initial_lookahead = MAX_MATCH + 1, exactly one literal will be emitted
 * before returning need_more.
 */
void test_deflate_rle_short_run_treated_as_literal(void) {
    z_stream strm;
    unsigned char out[32768];
    deflate_state *s = init_deflate_state(&strm, Z_RLE, out, sizeof(out));

    const unsigned initial_strstart = 1;
    const unsigned initial_lookahead = MAX_MATCH + 1;

    /* prev (index 0) = 'A', then exactly two 'A's, then a different byte, to avoid 3-match gate */
    s->window[0] = (Byte)'A';
    s->window[1] = (Byte)'A';
    s->window[2] = (Byte)'A';
    s->window[3] = (Byte)'B';
    /* Fill remainder with non-repeating sequence to avoid accidental runs */
    for (unsigned i = 4; i <= initial_strstart + initial_lookahead; i++) {
        s->window[i] = (Byte)(i & 0xFF);
    }

    s->strstart = initial_strstart;
    s->block_start = 0L;
    s->lookahead = initial_lookahead;
    s->match_length = 0;
    s->sym_next = 0;

    int r = test_deflate_rle(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(0, r); /* need_more */

    /* Exactly one literal emitted, so strstart advanced by 1, lookahead reduced to MAX_MATCH */
    TEST_ASSERT_EQUAL_UINT(initial_strstart + 1, s->strstart);
    TEST_ASSERT_EQUAL_UINT(MAX_MATCH, s->lookahead);
    TEST_ASSERT_TRUE(s->sym_next > 0);

    deflateEnd(&strm);
}

/* Test: With no lookahead and Z_NO_FLUSH, the function should return need_more
 * and not change state (no progress).
 */
void test_deflate_rle_no_input_no_flush_returns_need_more(void) {
    z_stream strm;
    unsigned char out[1024];
    deflate_state *s = init_deflate_state(&strm, Z_RLE, out, sizeof(out));

    /* Ensure no input/lookahead */
    s->strstart = 0;
    s->block_start = 0L;
    s->lookahead = 0;
    s->sym_next = 0;

    int r = test_deflate_rle(s, Z_NO_FLUSH);
    TEST_ASSERT_EQUAL_INT(0, r); /* need_more */
    TEST_ASSERT_EQUAL_UINT(0, s->strstart);
    TEST_ASSERT_EQUAL_UINT(0, s->lookahead);
    TEST_ASSERT_EQUAL_UINT(0, s->sym_next);

    deflateEnd(&strm);
}

/* Test: With no lookahead and Z_FINISH, the function should flush and return finish_done (3). */
void test_deflate_rle_finish_done_when_no_lookahead(void) {
    z_stream strm;
    unsigned char out[32768];
    deflate_state *s = init_deflate_state(&strm, Z_RLE, out, sizeof(out));

    s->strstart = 0;
    s->block_start = 0L;
    s->lookahead = 0;
    s->sym_next = 0;

    int r = test_deflate_rle(s, Z_FINISH);
    TEST_ASSERT_EQUAL_INT(3, r); /* finish_done */
    TEST_ASSERT_EQUAL_UINT(0, s->insert);

    deflateEnd(&strm);
}

int main(void) {
    UNITY_BEGIN();

    RUN_TEST(test_deflate_rle_consumes_long_run_and_returns_need_more);
    RUN_TEST(test_deflate_rle_emits_literals_until_need_more);
    RUN_TEST(test_deflate_rle_short_run_treated_as_literal);
    RUN_TEST(test_deflate_rle_no_input_no_flush_returns_need_more);
    RUN_TEST(test_deflate_rle_finish_done_when_no_lookahead);

    return UNITY_END();
}