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

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

/* Helper: decode a full Z85 buffer in one shot and verify result */
static void helper_decode_all(const char *enc, idx_t enc_len,
                              unsigned char *out, idx_t *out_len)
{
  struct base_decode_context ctx;
  z85_decode_ctx_init(&ctx);

  idx_t produced = 0;
  bool ok = z85_decode_ctx(&ctx, enc, enc_len, (char*)out, &produced);
  TEST_ASSERT_TRUE_MESSAGE(ok, "z85_decode_ctx failed on full buffer");

  /* Flush (inlen == 0) should succeed with no extra output. */
  idx_t flushn = 0;
  ok = z85_decode_ctx(&ctx, "", 0, (char*)out, &flushn);
  TEST_ASSERT_TRUE_MESSAGE(ok, "z85_decode_ctx flush failed");
  TEST_ASSERT_EQUAL_UINT64_MESSAGE(0, (uint64_t)flushn, "Unexpected flush output");

  *out_len = produced;
}

/* Test: empty input should produce no output and not crash */
void test_z85_encode_empty(void)
{
  char outbuf[1] = { (char)0xAB };
  idx_t outlen = z85_length(0);
  TEST_ASSERT_EQUAL_UINT64(0, (uint64_t)outlen);
  z85_encode((const char*)"", 0, outbuf, outlen);
  /* Nothing else to assert; absence of crash is success. */
  TEST_PASS();
}

/* Test: known vector 0x86 0x4F 0xD2 0x6F -> "Hello" */
void test_z85_encode_known_hello(void)
{
  const unsigned char in[4] = { 0x86, 0x4F, 0xD2, 0x6F };
  idx_t outlen = z85_length(4);
  TEST_ASSERT_EQUAL_UINT64(5, (uint64_t)outlen);
  char out[5];
  z85_encode((const char*)in, 4, out, outlen);
  const char expected[5] = { 'H','e','l','l','o' };
  TEST_ASSERT_EQUAL_MEMORY(expected, out, 5);
}

/* Test: known vector 8 bytes -> "HelloWorld" */
void test_z85_encode_known_helloworld(void)
{
  const unsigned char in[8] = {
    0x86, 0x4F, 0xD2, 0x6F,  /* "Hello" */
    0xB5, 0x59, 0xF7, 0x5B   /* "World" */
  };
  idx_t outlen = z85_length(8);
  TEST_ASSERT_EQUAL_UINT64(10, (uint64_t)outlen);
  char out[10];
  z85_encode((const char*)in, 8, out, outlen);
  const char expected[10] = {
    'H','e','l','l','o','W','o','r','l','d'
  };
  TEST_ASSERT_EQUAL_MEMORY(expected, out, 10);
}

/* Test: all zero input -> should encode to "00000" */
void test_z85_encode_all_zero_block(void)
{
  const unsigned char in[4] = { 0x00, 0x00, 0x00, 0x00 };
  idx_t outlen = z85_length(4);
  TEST_ASSERT_EQUAL_UINT64(5, (uint64_t)outlen);
  char out[5];
  z85_encode((const char*)in, 4, out, outlen);
  const char expected[5] = { '0','0','0','0','0' };
  TEST_ASSERT_EQUAL_MEMORY(expected, out, 5);
}

/* Test: all 0xFF input -> expected from manual computation "%nSc0" */
void test_z85_encode_all_ff_block(void)
{
  const unsigned char in[4] = { 0xFF, 0xFF, 0xFF, 0xFF };
  idx_t outlen = z85_length(4);
  TEST_ASSERT_EQUAL_UINT64(5, (uint64_t)outlen);
  char out[5];
  z85_encode((const char*)in, 4, out, outlen);
  const char expected[5] = { '%','n','S','c','0' };
  TEST_ASSERT_EQUAL_MEMORY(expected, out, 5);
}

/* Test: round-trip encode -> decode for various inputs */
void test_z85_encode_roundtrip_various(void)
{
  const unsigned char v1[4] = { 0x00, 0x01, 0x02, 0x03 };
  const unsigned char v2[4] = { 0xFF, 0xEE, 0xDD, 0xCC };
  const unsigned char v3[8] = { 0x00, 0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70 };
  const unsigned char v4[12] = {
    0x12,0x34,0x56,0x78, 0x9A,0xBC,0xDE,0xF0, 0x11,0x22,0x33,0x44
  };

  const struct { const unsigned char *in; idx_t inlen; } cases[] = {
    { v1, 4 }, { v2, 4 }, { v3, 8 }, { v4, 12 }
  };

  for (size_t i = 0; i < sizeof(cases)/sizeof(cases[0]); i++) {
    const unsigned char *in = cases[i].in;
    idx_t inlen = cases[i].inlen;
    idx_t enc_len = z85_length(inlen);

    char *enc = (char*)malloc((size_t)enc_len);
    TEST_ASSERT_NOT_NULL(enc);

    z85_encode((const char*)in, inlen, enc, enc_len);

    unsigned char *dec = (unsigned char*)malloc((size_t)inlen);
    TEST_ASSERT_NOT_NULL(dec);
    idx_t produced = 0;
    helper_decode_all(enc, enc_len, dec, &produced);

    TEST_ASSERT_EQUAL_UINT64((uint64_t)inlen, (uint64_t)produced);
    TEST_ASSERT_EQUAL_MEMORY(in, dec, (size_t)inlen);

    free(enc);
    free(dec);
  }
}

/* Test: ensure no overwrite occurs beyond provided outlen */
void test_z85_encode_no_overwrite_beyond_outlen(void)
{
  const unsigned char in[4] = { 0x86, 0x4F, 0xD2, 0x6F }; /* "Hello" */
  idx_t outlen = z85_length(4); /* 5 */

  size_t bufsize = (size_t)outlen + 8; /* extra padding */
  unsigned char *buf = (unsigned char*)malloc(bufsize);
  TEST_ASSERT_NOT_NULL(buf);
  memset(buf, 0xAA, bufsize);

  z85_encode((const char*)in, 4, (char*)buf, outlen);

  const unsigned char expected[5] = { 'H','e','l','l','o' };
  TEST_ASSERT_EQUAL_MEMORY(expected, buf, 5);
  for (size_t i = (size_t)outlen; i < bufsize; i++) {
    TEST_ASSERT_EQUAL_HEX8_MESSAGE(0xAA, buf[i], "Buffer overwrite beyond outlen detected");
  }

  free(buf);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_z85_encode_empty);
  RUN_TEST(test_z85_encode_known_hello);
  RUN_TEST(test_z85_encode_known_helloworld);
  RUN_TEST(test_z85_encode_all_zero_block);
  RUN_TEST(test_z85_encode_all_ff_block);
  RUN_TEST(test_z85_encode_roundtrip_various);
  RUN_TEST(test_z85_encode_no_overwrite_beyond_outlen);
  return UNITY_END();
}