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

/* setUp and tearDown required by Unity */
void setUp(void) {
  /* No global setup needed */
}
void tearDown(void) {
  /* No global teardown needed */
}

/* Helper: compare encoded output to expected C string literal */
static void assert_encoded_equals_literal(const char *input, size_t inlen, const char *expected_literal)
{
  idx_t outlen = BASE64_LENGTH((idx_t)inlen);
  size_t expected_len = strlen(expected_literal);
  TEST_ASSERT_EQUAL_size_t(expected_len, (size_t)outlen);

  char *out = (char *)malloc(outlen ? outlen : 1);
  TEST_ASSERT_NOT_NULL(out);

  base64url_encode(input, (idx_t)inlen, out, outlen);

  TEST_ASSERT_EQUAL_MEMORY(expected_literal, out, outlen);

  free(out);
}

/* Test 1: Empty input should produce empty output and not write anything */
static void test_base64url_encode_empty(void)
{
  const char *in = "";
  idx_t outlen = BASE64_LENGTH((idx_t)0);
  TEST_ASSERT_EQUAL_INT(0, (int)outlen);

  char sentinel = (char)0xA5;
  char *out = &sentinel;

  base64url_encode(in, 0, out, outlen);

  /* Since outlen==0, output buffer must remain untouched */
  TEST_ASSERT_EQUAL_HEX8(0xA5, (unsigned char)sentinel);
}

/* Test 2: RFC 4648 basic vectors (no URL-specific changes needed) */
static void test_base64url_encode_rfc4648_vectors(void)
{
  assert_encoded_equals_literal("f",    1, "Zg==");
  assert_encoded_equals_literal("fo",   2, "Zm8=");
  assert_encoded_equals_literal("foo",  3, "Zm9v");
  assert_encoded_equals_literal("foob", 4, "Zm9vYg==");
  assert_encoded_equals_literal("fooba",5, "Zm9vYmE=");
  assert_encoded_equals_literal("foobar",6,"Zm9vYmFy");
}

/* Test 3: Verify translation of '+' -> '-' and '/' -> '_' */
static void test_base64url_encode_translates_plus_and_slash(void)
{
  /* Crafted 3-byte input producing "+/+A" in standard base64:
     0xFB, 0xFF, 0x80
     g0=62('+'), g1=63('/'), g2=62('+'), g3=0('A') */
  unsigned char in[3] = { 0xFB, 0xFF, 0x80 };
  idx_t outlen = BASE64_LENGTH((idx_t)sizeof(in));
  TEST_ASSERT_EQUAL_INT(4, (int)outlen);

  char *out = (char *)malloc(outlen);
  TEST_ASSERT_NOT_NULL(out);

  base64url_encode((const char *)in, (idx_t)sizeof(in), out, outlen);

  const char expected[4] = { '-', '_', '-', 'A' };
  TEST_ASSERT_EQUAL_MEMORY(expected, out, 4);

  free(out);
}

/* Test 4: Repeated pattern to ensure replacement across the entire buffer */
static void test_base64url_encode_repeated_pattern(void)
{
  /* Repeat the 3-byte pattern {0xFB,0xFF,0x80} 10 times */
  enum { REPS = 10 };
  unsigned char in[3 * REPS];
  for (int i = 0; i < REPS; i++) {
    in[3*i + 0] = 0xFB;
    in[3*i + 1] = 0xFF;
    in[3*i + 2] = 0x80;
  }

  idx_t outlen = BASE64_LENGTH((idx_t)sizeof(in));
  TEST_ASSERT_EQUAL_INT(4 * REPS, (int)outlen);

  char *out = (char *)malloc(outlen);
  TEST_ASSERT_NOT_NULL(out);

  base64url_encode((const char *)in, (idx_t)sizeof(in), out, outlen);

  /* Expected is "-_-A" repeated REPS times */
  char expected[4 * REPS];
  for (int i = 0; i < REPS; i++) {
    expected[4*i + 0] = '-';
    expected[4*i + 1] = '_';
    expected[4*i + 2] = '-';
    expected[4*i + 3] = 'A';
  }

  TEST_ASSERT_EQUAL_MEMORY(expected, out, outlen);

  free(out);
}

/* Test 5: For a broader input range, compare to base64 + manual URL translation */
static void test_base64url_encode_matches_base64_then_translate(void)
{
  /* Use a variety of byte values */
  enum { LEN = 256 };
  unsigned char in[LEN];
  for (int i = 0; i < LEN; i++)
    in[i] = (unsigned char)i;

  idx_t outlen = BASE64_LENGTH((idx_t)LEN);

  char *got = (char *)malloc(outlen);
  char *ref = (char *)malloc(outlen);
  TEST_ASSERT_NOT_NULL(got);
  TEST_ASSERT_NOT_NULL(ref);

  /* Compute reference: base64 then translate '+'->'-', '/'->'_' */
  base64_encode((const char *)in, (idx_t)LEN, ref, outlen);
  for (idx_t i = 0; i < outlen; i++) {
    if (ref[i] == '+') ref[i] = '-';
    else if (ref[i] == '/') ref[i] = '_';
  }

  base64url_encode((const char *)in, (idx_t)LEN, got, outlen);

  TEST_ASSERT_EQUAL_MEMORY(ref, got, outlen);

  free(got);
  free(ref);
}

int main(void)
{
  UNITY_BEGIN();
  RUN_TEST(test_base64url_encode_empty);
  RUN_TEST(test_base64url_encode_rfc4648_vectors);
  RUN_TEST(test_base64url_encode_translates_plus_and_slash);
  RUN_TEST(test_base64url_encode_repeated_pattern);
  RUN_TEST(test_base64url_encode_matches_base64_then_translate);
  return UNITY_END();
}