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

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

/* Helper to compute the expected value, matching the implementation comment */
static idx_t expected_base58_length(idx_t len)
{
  return ((len + 99) / 100) * 138 + 1;
}

/* Helper to fill buffer deterministically */
static void fill_buffer(char *buf, size_t len)
{
  for (size_t i = 0; i < len; i++)
    buf[i] = (char)((i * 37u + 13u) & 0xFF);
}

/* Test: len == 0 and a few basic values */
static void test_base58_length_zero_and_basics(void)
{
  TEST_ASSERT_EQUAL_INT64((long long)1, (long long)base58_length(0));

  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(1),   (long long)base58_length(1));
  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(50),  (long long)base58_length(50));
  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(99),  (long long)base58_length(99));
  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(100), (long long)base58_length(100));
  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(101), (long long)base58_length(101));
  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(200), (long long)base58_length(200));
  TEST_ASSERT_EQUAL_INT64((long long)expected_base58_length(201), (long long)base58_length(201));
}

/* Test: boundaries and step increments */
static void test_base58_length_boundaries_and_steps(void)
{
  /* From 1..100 should be constant (== 138+1) */
  idx_t first_block = base58_length(1);
  for (idx_t i = 1; i <= 100; i++)
    TEST_ASSERT_EQUAL_INT64((long long)first_block, (long long)base58_length(i));

  /* 101..200 should be exactly +138 from previous block */
  idx_t second_block = base58_length(101);
  TEST_ASSERT_EQUAL_INT64((long long)(first_block + 138), (long long)second_block);
  for (idx_t i = 101; i <= 200; i++)
    TEST_ASSERT_EQUAL_INT64((long long)second_block, (long long)base58_length(i));

  /* Check next boundary as well */
  idx_t third_block = base58_length(201);
  TEST_ASSERT_EQUAL_INT64((long long)(second_block + 138), (long long)third_block);
}

/* Test: positivity for several values, including 0 */
static void test_base58_length_positive(void)
{
  TEST_ASSERT_TRUE(base58_length(0) > 0);
  TEST_ASSERT_TRUE(base58_length(1) > 0);
  TEST_ASSERT_TRUE(base58_length(50) > 0);
  TEST_ASSERT_TRUE(base58_length(100) > 0);
  TEST_ASSERT_TRUE(base58_length(1000) > 0);
}

/* Functional test: ensure base58_length(len) is sufficient for encoding random data */
static void test_base58_length_sufficient_for_encoding_various_sizes(void)
{
#if BASE_TYPE == 42
  /* We will call base58_encode, which asserts base_length(...) <= outcap. */
  idx_t (*saved_base_length)(idx_t) = base_length;
  base_length = base58_length;

  size_t sizes[] = {0, 1, 2, 10, 32, 50, 99, 100, 101, 256, 1000};
  for (size_t si = 0; si < sizeof(sizes)/sizeof(sizes[0]); si++)
    {
      size_t inlen = sizes[si];
      char *inbuf = NULL;
      if (inlen)
        {
          inbuf = (char *)malloc(inlen);
          TEST_ASSERT_NOT_NULL(inbuf);
          fill_buffer(inbuf, inlen);
        }

      idx_t outcap = base58_length((idx_t)inlen);
      char *outbuf = (char *)malloc((size_t)outcap);
      TEST_ASSERT_NOT_NULL(outbuf);

      idx_t outlen = outcap;
      base58_encode(inbuf ? inbuf : (char *)"", inlen, outbuf, &outlen);

      /* Must fit within computed capacity */
      TEST_ASSERT_TRUE(outlen <= outcap);

      free(outbuf);
      free(inbuf);
    }

  base_length = saved_base_length;
#else
  /* This test is only relevant for basenc (BASE_TYPE==42) where base58 is present. */
  TEST_IGNORE_MESSAGE("Functional base58 encode sufficiency test skipped (BASE_TYPE != 42)");
#endif
}

/* Functional test: all-zero input encodes to that many '1's, ensure capacity covers it */
static void test_base58_length_all_zero_input(void)
{
#if BASE_TYPE == 42
  idx_t (*saved_base_length)(idx_t) = base_length;
  base_length = base58_length;

  size_t inlen = 150; /* spans across the first two step ranges */
  char *inbuf = (char *)calloc(inlen, 1); /* all zeros */
  TEST_ASSERT_NOT_NULL(inbuf);

  idx_t outcap = base58_length((idx_t)inlen);
  char *outbuf = (char *)malloc((size_t)outcap);
  TEST_ASSERT_NOT_NULL(outbuf);

  idx_t outlen = outcap;
  base58_encode(inbuf, inlen, outbuf, &outlen);

  /* For all-zero input, output is exactly 'inlen' copies of '1' */
  TEST_ASSERT_EQUAL_UINT64((unsigned long long)inlen, (unsigned long long)outlen);
  for (idx_t i = 0; i < outlen; i++)
    TEST_ASSERT_EQUAL_INT('1', (unsigned char)outbuf[i]);

  free(outbuf);
  free(inbuf);
  base_length = saved_base_length;
#else
  TEST_IGNORE_MESSAGE("Zero-input encoding test skipped (BASE_TYPE != 42)");
#endif
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_base58_length_zero_and_basics);
  RUN_TEST(test_base58_length_boundaries_and_steps);
  RUN_TEST(test_base58_length_positive);
  RUN_TEST(test_base58_length_sufficient_for_encoding_various_sizes);
  RUN_TEST(test_base58_length_all_zero_input);
  return UNITY_END();
}