// b64.c – base64url encoding/decoding with arena allocation

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-function"

static const char *base64_url_alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
static const char *base64_std_alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

// base64_encode main function, called by others with different alphabet/padding rules
// it does NOT add a nul char at the end of output
static void base64_encode_generic(char const* alphabet, bool with_padding, byte *dst, int dst_len, byte const *src, int src_len) {
  int i = 0, j = 0;
  unsigned char a, b, c;

  while (src_len >= 3) {
    a = src[i++];
    b = src[i++];
    c = src[i++];
    dst[j++] = alphabet[a >> 2];
    dst[j++] = alphabet[((a & 0x03) << 4) | (b >> 4)];
    dst[j++] = alphabet[((b & 0x0f) << 2) | (c >> 6)];
    dst[j++] = alphabet[c & 0x3f];
    src_len -= 3;
  }

  if (src_len == 1) {
    a = src[i++];
    dst[j++] = alphabet[a >> 2];
    dst[j++] = alphabet[(a & 0x03) << 4];
    if(with_padding){
      dst[j++] = '=';
      dst[j++] = '=';
    }
  } else if (src_len == 2) {
    a = src[i++];
    b = src[i++];
    dst[j++] = alphabet[a >> 2];
    dst[j++] = alphabet[((a & 0x03) << 4) | (b >> 4)];
    dst[j++] = alphabet[(b & 0x0f) << 2];
    if(with_padding){
      dst[j++] = '=';
    }
  }
  assert(j<=dst_len);
}

static void base64_encode_std(byte *dst, int dst_len, byte const *src, int src_len) {
  enum{with_padding=true};
  base64_encode_generic(base64_std_alphabet, with_padding, dst, dst_len, src, src_len);
}

static int b64url_char_index(char c) {
    if (c >= 'A' && c <= 'Z') return c - 'A';
    if (c >= 'a' && c <= 'z') return 26 + (c - 'a');
    if (c >= '0' && c <= '9') return 52 + (c - '0');
    if (c == '-') return 62;
    if (c == '_') return 63;
    return -1;
}

/* ---------------------------------------------------------------------------
   Encode plain bytes to base64url (no padding).
   Returns a str allocated in arena A, or an empty str on error.
   --------------------------------------------------------------------------- */
static str base64url_encode(arena *A, str plain) {
  if (plain.len == 0) return (str){0};
  // Calculate output length without padding
  size_t out_len = 4 * (plain.len / 3);
  size_t rem = plain.len % 3;
  if (rem == 1) out_len += 2;
  else if (rem == 2) out_len += 3;

  byte *out = new_array(A, byte, out_len);
  if (!out) return (str){0};

  enum{no_padding=false};
  base64_encode_generic(base64_url_alphabet, no_padding, out, out_len, plain.data, plain.len);

  return str_from_buf(out, out_len);
}

/* ---------------------------------------------------------------------------
   Decode base64url string to raw bytes.
   Handles input with or without padding.
   Returns a str allocated in arena A, or an empty str on error.
   --------------------------------------------------------------------------- */
static str base64url_decode(arena *A, str b64encoded) {
    if (b64encoded.len == 0) return (str){0};
    const char *in = (const char*)b64encoded.data;
    size_t in_len = (size_t)b64encoded.len;

    // Ignore any trailing padding '=' characters
    while (in_len > 0 && in[in_len-1] == '=') in_len--;

    // Each group of 4 chars becomes 3 bytes, except last group may be shorter
    size_t out_len = (in_len * 3) / 4;
    uint8_t *out = new_array(A, uint8_t, out_len);
    if (!out) return (str){0};

    size_t i = 0, j = 0;
    uint32_t buffer = 0;
    int bits = 0;

    for (i = 0; i < in_len; i++) {
        int idx = b64url_char_index(in[i]);
        if (idx < 0) return (str){0}; // invalid character
        buffer = (buffer << 6) | idx;
        bits += 6;
        if (bits >= 8) {
            bits -= 8;
            out[j++] = (buffer >> bits) & 0xFF;
            if (j >= out_len) break;
        }
    }
    return str_from_buf(out, j);
}

#pragma GCC diagnostic pop