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

/* The target function and helpers are in the same translation unit. */
/* We can reference base58_decode and base58_to_gmp directly. */

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

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

/* Helper: Convert Base58 alphabet string to the GMP digit string that
   base58_decode expects, using the internal base58_to_gmp table. */
static char *alphabet_to_gmp_digits(const char *alphabet, size_t *out_len)
{
  if (!alphabet) return NULL;
  size_t n = strlen(alphabet);
  char *digits = (char *)malloc(n + 1);
  TEST_ASSERT_NOT_NULL(digits);
  for (size_t i = 0; i < n; i++) {
    signed char m = base58_to_gmp[(unsigned char)alphabet[i]];
    TEST_ASSERT_MESSAGE(m != -1, "Invalid Base58 alphabet character in test input");
    digits[i] = (char)m;
  }
  digits[n] = '\0';
  if (out_len) *out_len = n;
  return digits;
}

/* Helper: Build a GMP base-58 digit string from raw bytes as base58_decode
   expects: '0' chars for leading zero bytes, followed by mpz_get_str(base=58)
   representation of the remaining number. Returns malloc'd string and length. */
static char *bytes_to_gmp58_digits(const unsigned char *data, size_t len, size_t *digits_len)
{
  size_t ones = 0;
  while (ones < len && data[ones] == 0) ones++;

  if (len - ones == 0) {
    /* All zeros: just ones copies of '0'. */
    char *digits = (char *)malloc(ones + 1);
    TEST_ASSERT_NOT_NULL(digits);
    memset(digits, '0', ones);
    digits[ones] = '\0';
    if (digits_len) *digits_len = ones;
    return digits;
  }

  mpz_t num;
  mpz_init(num);
  mpz_import(num, len - ones, 1, 1, 0, 0, data + ones);

  /* Determine needed capacity: ones + mpz_sizeinbase(num, 58) + 2 for safety. */
  size_t cap = ones + mpz_sizeinbase(num, 58) + 2;
  char *digits = (char *)malloc(cap);
  TEST_ASSERT_NOT_NULL(digits);
  memset(digits, '0', ones);
  char *p = digits + ones;
  char *ret = mpz_get_str(p, 58, num);
  TEST_ASSERT_EQUAL_PTR(p, ret);
  size_t tail_len = strlen(p);
  if (digits_len) *digits_len = ones + tail_len;
  mpz_clear(num);
  return digits;
}

/* Test: empty input should succeed and produce no output. */
void test_base58_decode_empty(void)
{
  char outbuf[1];
  idx_t outlen = 0; /* Precondition: data_len <= outlen (0 <= 0) holds */
  bool ok = base58_decode("", 0, outbuf, &outlen);
  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(0, (unsigned long long)outlen);
}

/* Test: only leading zeros (Base58 alphabet "1" -> GMP digit '0'). */
void test_base58_decode_only_ones_single(void)
{
  size_t inlen = 0;
  char *gmp = alphabet_to_gmp_digits("1", &inlen); /* becomes "0" */
  TEST_ASSERT_EQUAL_size_t(1, inlen);

  unsigned char outbuf[4] = {0xAA, 0xAA, 0xAA, 0xAA};
  idx_t outlen = (idx_t)inlen; /* satisfy precondition */
  bool ok = base58_decode(gmp, inlen, (char *)outbuf, &outlen);
  free(gmp);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(1, (unsigned long long)outlen);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[0]);
}

/* Test: three leading '1's then '2' (alphabet) -> GMP "0001" -> bytes 00 00 00 01. */
void test_base58_decode_ones_plus_digit(void)
{
  size_t inlen = 0;
  char *gmp = alphabet_to_gmp_digits("1112", &inlen); /* becomes "0001" */

  unsigned char outbuf[16];
  memset(outbuf, 0xCC, sizeof(outbuf));
  idx_t outlen = (idx_t)inlen; /* satisfy precondition */
  bool ok = base58_decode(gmp, inlen, (char *)outbuf, &outlen);
  free(gmp);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(4, (unsigned long long)outlen);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[0]);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[1]);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[2]);
  TEST_ASSERT_EQUAL_UINT8(0x01, outbuf[3]);
}

/* Test: roundtrip for small arbitrary bytes without leading zeros. */
void test_base58_decode_roundtrip_small(void)
{
  unsigned char input[] = { 0x01, 0x02, 0xFF };
  size_t digits_len = 0;
  char *digits = bytes_to_gmp58_digits(input, sizeof(input), &digits_len);

  /* Output buffer longer than input digits to check function sets exact outlen. */
  size_t outcap = digits_len + 8;
  unsigned char *outbuf = (unsigned char *)malloc(outcap);
  TEST_ASSERT_NOT_NULL(outbuf);
  idx_t outlen = (idx_t)outcap; /* satisfy precondition (>= data_len) */

  bool ok = base58_decode(digits, digits_len, (char *)outbuf, &outlen);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(sizeof(input), (unsigned long long)outlen);
  TEST_ASSERT_EQUAL_UINT8_ARRAY(input, outbuf, sizeof(input));

  free(digits);
  free(outbuf);
}

/* Test: roundtrip with leading zeros preserved. */
void test_base58_decode_roundtrip_leading_zeros(void)
{
  unsigned char input[] = { 0x00, 0x00, 0x00, 0x01, 0x23, 0x45 };
  size_t digits_len = 0;
  char *digits = bytes_to_gmp58_digits(input, sizeof(input), &digits_len);

  unsigned char *outbuf = (unsigned char *)malloc(digits_len + 16);
  TEST_ASSERT_NOT_NULL(outbuf);
  idx_t outlen = (idx_t)(digits_len + 16); /* precondition */

  bool ok = base58_decode(digits, digits_len, (char *)outbuf, &outlen);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(sizeof(input), (unsigned long long)outlen);
  TEST_ASSERT_EQUAL_UINT8_ARRAY(input, outbuf, sizeof(input));

  free(digits);
  free(outbuf);
}

/* Test: invalid digit for base 58 (e.g., 'z' corresponds to value 61 in GMP digits). */
void test_base58_decode_invalid_digit(void)
{
  const char *digits = "0z"; /* one leading zero, then invalid digit for base 58 */
  unsigned char outbuf[8] = {0};
  idx_t outlen = 2; /* data_len (2) <= outlen (2) */

  bool ok = base58_decode(digits, 2, (char *)outbuf, &outlen);

  TEST_ASSERT_FALSE(ok);
  TEST_ASSERT_EQUAL_UINT64(0, (unsigned long long)outlen);
}

/* Test: ensure larger output buffer than input length works and outlen shrinks to actual size. */
void test_base58_decode_large_outbuf(void)
{
  size_t inlen = 0;
  char *gmp = alphabet_to_gmp_digits("1112", &inlen); /* GMP: "0001" */

  unsigned char outbuf[32];
  idx_t outlen = 10; /* >= data_len (4) satisfies precondition */

  bool ok = base58_decode(gmp, inlen, (char *)outbuf, &outlen);
  free(gmp);

  TEST_ASSERT_TRUE(ok);
  TEST_ASSERT_EQUAL_UINT64(4, (unsigned long long)outlen);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[0]);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[1]);
  TEST_ASSERT_EQUAL_UINT8(0x00, outbuf[2]);
  TEST_ASSERT_EQUAL_UINT8(0x01, outbuf[3]);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_base58_decode_empty);
  RUN_TEST(test_base58_decode_only_ones_single);
  RUN_TEST(test_base58_decode_ones_plus_digit);
  RUN_TEST(test_base58_decode_roundtrip_small);
  RUN_TEST(test_base58_decode_roundtrip_leading_zeros);
  RUN_TEST(test_base58_decode_invalid_digit);
  RUN_TEST(test_base58_decode_large_outbuf);
  return UNITY_END();
}