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

/* External wrapper generated in the module for testing the local function */
void test_flush_pending(z_streamp strm);

/* Global stream reused per test */
static z_stream strm;

void setUp(void) {
    memset(&strm, 0, sizeof(strm));
    int ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    TEST_ASSERT_EQUAL_INT_MESSAGE(Z_OK, ret, "deflateInit failed");
}

void tearDown(void) {
    /* Only call deflateEnd if state is initialized */
    if (strm.state != Z_NULL) {
        int ret = deflateEnd(&strm);
        TEST_ASSERT_TRUE_MESSAGE(ret == Z_OK || ret == Z_DATA_ERROR, "deflateEnd unexpected status");
        memset(&strm, 0, sizeof(strm));
    }
}

/* Helper: produce a partially flushed zlib header to ensure pending > 0.
   Writes one byte of the header out (to out_first[0]) and leaves the rest pending.
   Returns number of pending bytes after the call. */
static unsigned make_header_with_one_pending(unsigned char *out_first) {
    unsigned pending = 0;
    int bits = 0;

    unsigned char small_out[1];
    strm.next_in = Z_NULL;
    strm.avail_in = 0;
    strm.next_out = small_out;
    strm.avail_out = sizeof(small_out);

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

    if (out_first) {
        *out_first = small_out[0];
    }

    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pending, &bits));
    TEST_ASSERT_MESSAGE(pending > 0, "Expected at least one pending byte after partial header flush");
    return pending;
}

/* Test: flush_pending flushes remaining header bytes to next_out */
void test_flush_pending_flushes_remaining_header(void) {
    unsigned char first_byte = 0;
    unsigned pending_before = make_header_with_one_pending(&first_byte);

    /* Prepare output buffer to receive remaining pending bytes */
    unsigned char out_rest[8];
    memset(out_rest, 0, sizeof(out_rest));
    unsigned char *next_out_before = out_rest;

    /* Set stream output to our buffer */
    strm.next_out = out_rest;
    strm.avail_out = (uInt)sizeof(out_rest);

    uLong total_out_before = strm.total_out;
    uInt avail_out_before = strm.avail_out;

    /* Verify pending exists before calling the function under test */
    unsigned pchk = 0; int bits = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pchk, &bits));
    TEST_ASSERT_EQUAL_UINT(pending_before, pchk);

    /* Call the wrapper for local flush_pending */
    test_flush_pending(&strm);

    /* After flush, pending should be zero */
    unsigned pending_after = 0; int bits_after = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pending_after, &bits_after));
    TEST_ASSERT_EQUAL_UINT_MESSAGE(0, pending_after, "Pending should be zero after flush_pending");

    /* Check counters and pointer movement */
    uLong total_out_after = strm.total_out;
    uInt avail_out_after = strm.avail_out;

    TEST_ASSERT_EQUAL_UINT_MESSAGE(pending_before, (unsigned)(total_out_after - total_out_before),
                                   "total_out should increase by flushed pending amount");
    TEST_ASSERT_EQUAL_UINT_MESSAGE(pending_before, (unsigned)(avail_out_before - avail_out_after),
                                   "avail_out should decrease by flushed pending amount");
    TEST_ASSERT_EQUAL_PTR_MESSAGE(next_out_before + pending_before, strm.next_out,
                                  "next_out should advance by flushed pending amount");

    /* Sanity-check the header bytes formed: (first_byte, second_byte) */
    /* For the standard zlib header, the 16-bit value is divisible by 31,
       and CM (lower 4 bits of first byte) must be Z_DEFLATED (8). */
    TEST_ASSERT_TRUE_MESSAGE(pending_before >= 1, "Expected at least one remaining header byte");
    uint16_t header = ((uint16_t)first_byte << 8) | out_rest[0];
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, header % 31, "Zlib header check (mod 31) failed");
    TEST_ASSERT_EQUAL_INT_MESSAGE(Z_DEFLATED, (first_byte & 0x0F), "Zlib header CM not deflate (8)");
}

/* Test: when no pending data, flush_pending is a no-op */
void test_flush_pending_no_pending_noop(void) {
    /* Initially after deflateInit, pending should be zero */
    unsigned pending = 0; int bits = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pending, &bits));
    TEST_ASSERT_EQUAL_UINT(0, pending);

    unsigned char outbuf[4] = {0};
    strm.next_out = outbuf;
    strm.avail_out = sizeof(outbuf);
    uLong total_out_before = strm.total_out;
    uInt avail_out_before = strm.avail_out;
    unsigned char *next_out_before = strm.next_out;

    test_flush_pending(&strm);

    /* Nothing should have changed */
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pending, &bits));
    TEST_ASSERT_EQUAL_UINT(0, pending);
    TEST_ASSERT_EQUAL_UINT(total_out_before, strm.total_out);
    TEST_ASSERT_EQUAL_UINT(avail_out_before, strm.avail_out);
    TEST_ASSERT_EQUAL_PTR(next_out_before, strm.next_out);
}

/* Test: when avail_out is zero, flush_pending does nothing; after providing space, it flushes */
void test_flush_pending_zero_avail_out_then_flush(void) {
    unsigned char first_byte = 0;
    unsigned pending_before = make_header_with_one_pending(&first_byte);

    /* First call with no output space */
    unsigned char outbuf[4] = {0};
    strm.next_out = outbuf;
    strm.avail_out = 0;

    uLong total_out_before = strm.total_out;

    test_flush_pending(&strm);

    /* Verify nothing changed */
    unsigned pending_mid = 0; int bits_mid = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pending_mid, &bits_mid));
    TEST_ASSERT_EQUAL_UINT_MESSAGE(pending_before, pending_mid, "Pending should be unchanged with avail_out == 0");
    TEST_ASSERT_EQUAL_UINT_MESSAGE(total_out_before, strm.total_out, "total_out should be unchanged with avail_out == 0");

    /* Now provide space and flush */
    strm.avail_out = sizeof(outbuf);
    unsigned char *next_out_before = strm.next_out;

    test_flush_pending(&strm);

    unsigned pending_after = 0; int bits_after = 0;
    TEST_ASSERT_EQUAL_INT(Z_OK, deflatePending(&strm, &pending_after, &bits_after));
    TEST_ASSERT_EQUAL_UINT(0, pending_after);
    TEST_ASSERT_EQUAL_PTR(next_out_before + pending_mid, strm.next_out);
    TEST_ASSERT_EQUAL_UINT((uInt)pending_mid, (uInt)(sizeof(outbuf) - strm.avail_out));
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_flush_pending_flushes_remaining_header);
    RUN_TEST(test_flush_pending_no_pending_noop);
    RUN_TEST(test_flush_pending_zero_avail_out_then_flush);
    return UNITY_END();
}