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

/* The test file is included into the same translation unit as the
   basenc implementation, so we can access internal symbols directly. */

void setUp(void) {
  /* Ensure base58_encode uses the correct length function. */
  base_length = base58_length;
}

void tearDown(void) {
  /* nothing */
}

static void check_encode(const unsigned char *data, size_t data_len,
                         const char *expected)
{
  idx_t cap = base58_length((idx_t)data_len);
  char *out = (char *)malloc((size_t)cap);
  TEST_ASSERT_NOT_NULL(out);
  memset(out, 0xAA, (size_t)cap);

  idx_t outlen = cap;
  base58_encode((const char *)data, data_len, out, &outlen);

  size_t exp_len = strlen(expected);
  TEST_ASSERT_EQUAL_UINT64(exp_len, (uint64_t)outlen);
  if (exp_len > 0)
    TEST_ASSERT_EQUAL_INT(0, memcmp(out, expected, exp_len));
  /* Ensure no bytes beyond returned length were modified. */
  for (idx_t i = outlen; i < cap; i++)
    TEST_ASSERT_EQUAL_HEX8(0xAA, (unsigned char)out[i]);

  free(out);
}

static size_t count_leading_zero_bytes(const unsigned char *p, size_t n)
{
  size_t i = 0;
  while (i < n && p[i] == 0x00) i++;
  return i;
}

static void test_base58_encode_empty(void)
{
  check_encode(NULL, 0, "");
}

static void test_base58_encode_only_zeros(void)
{
  /* 1 to 4 leading zeros */
  unsigned char z1[] = {0x00};
  unsigned char z2[] = {0x00, 0x00};
  unsigned char z3[] = {0x00, 0x00, 0x00};
  unsigned char z4[] = {0x00, 0x00, 0x00, 0x00};

  check_encode(z1, sizeof z1, "1");
  check_encode(z2, sizeof z2, "11");
  check_encode(z3, sizeof z3, "111");
  check_encode(z4, sizeof z4, "1111");
}

static void test_base58_encode_single_bytes(void)
{
  unsigned char v1[] = {0x01};      /* value 1 -> "2" */
  unsigned char vff[] = {0xFF};     /* value 255 -> "5Q" */

  check_encode(v1, sizeof v1, "2");
  check_encode(vff, sizeof vff, "5Q");
}

static void test_base58_encode_prefix_zeros_then_value(void)
{
  unsigned char a[] = {0x00, 0x01};      /* 0x00,1 => "12" */
  unsigned char b[] = {0x00, 0x00, 0xFF};/* two zeros then 255 => "115Q" */

  check_encode(a, sizeof a, "12");
  check_encode(b, sizeof b, "115Q");
}

static void test_base58_encode_small_multibyte_examples(void)
{
  /* 0x0100 = 256 => base58 digits 4,24 => "5R" */
  unsigned char v256[] = {0x01, 0x00};
  check_encode(v256, sizeof v256, "5R");

  /* 0x01 0x02 0x03 = 66051 => digits 19,36,47 => "Ldp" */
  unsigned char v123[] = {0x01, 0x02, 0x03};
  check_encode(v123, sizeof v123, "Ldp");
}

static bool is_forbidden_base58_char(char c)
{
  /* Base58 alphabet excludes these visually ambiguous characters */
  return (c == '0' || c == 'O' || c == 'I' || c == 'l');
}

static void test_base58_encode_long_input_properties(void)
{
  /* Create a non-zero-leading buffer: 0x01, 0x02, ..., 0xFF, 0x00, 0x01 */
  size_t n = 300;
  unsigned char *buf = (unsigned char *)malloc(n);
  TEST_ASSERT_NOT_NULL(buf);
  for (size_t i = 0; i < n; i++)
    buf[i] = (unsigned char)((i + 1) & 0xFF);
  buf[0] = 0x01; /* ensure no leading zero */

  idx_t cap = base58_length((idx_t)n);
  char *out = (char *)malloc((size_t)cap);
  TEST_ASSERT_NOT_NULL(out);
  memset(out, 0xAA, (size_t)cap);

  idx_t outlen = cap;
  base58_encode((const char *)buf, n, out, &outlen);

  TEST_ASSERT_TRUE(outlen > 0);
  /* Since first byte != 0x00, output must not start with '1' (value 0 digit). */
  TEST_ASSERT_NOT_EQUAL('1', out[0]);

  /* Check that no forbidden characters appear in output. */
  for (idx_t i = 0; i < outlen; i++)
    TEST_ASSERT_FALSE(is_forbidden_base58_char(out[i]));

  /* Ensure no writes past returned length. */
  for (idx_t i = outlen; i < cap; i++)
    TEST_ASSERT_EQUAL_HEX8(0xAA, (unsigned char)out[i]);

  free(out);
  free(buf);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_base58_encode_empty);
  RUN_TEST(test_base58_encode_only_zeros);
  RUN_TEST(test_base58_encode_single_bytes);
  RUN_TEST(test_base58_encode_prefix_zeros_then_value);
  RUN_TEST(test_base58_encode_small_multibyte_examples);
  RUN_TEST(test_base58_encode_long_input_properties);
  return UNITY_END();
}