// vim:fileencoding=utf-8:foldmethod=marker

#include "base64.h"
#include "../wapp/wapp.h"
#include <math.h>

internal Str8RO alphabet = wapp_str8_lit_ro("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/");

internal c8 char_at(u8 index);
internal u8 char_index(c8 c);
internal u64 calc_encode_length(Str8RO *input);
internal u64 calc_decode_length(Str8RO *input);

Str8 *encode(Allocator *allocator, Str8RO *input) {
  if (input->size == 0) {
    return NULL;
  }

  u64 n_out = calc_encode_length(input);
  Str8 *out = wapp_str8_alloc_and_fill_buf(allocator, n_out);
  if (!out) return NULL;

  u8 buf[3] = {0};
  u8 count  = 0;
  u64 iout  = 0;


  for (u64 i = 0; i < input->size; ++i) {
    buf[count] = wapp_str8_get(input, i);
    count += 1;
    if (count == 3) {
      wapp_str8_set(out, iout, char_at(buf[0] >> 2));
      wapp_str8_set(out, iout + 1, char_at(((buf[0] & 0x03) << 4) + (buf[1] >> 4)));
      wapp_str8_set(out, iout + 2, char_at(((buf[1] & 0x0f) << 2) + (buf[2] >> 6)));
      wapp_str8_set(out, iout + 3, char_at(buf[2] & 0x3f));

      iout += 4;
      count = 0;
    }
  }

  if (count == 1)  {
    wapp_str8_set(out, iout, char_at(buf[0] >> 2));
    wapp_str8_set(out, iout + 1, char_at((buf[0] & 0x03) << 4));
    wapp_str8_set(out, iout + 2, '=');
    wapp_str8_set(out, iout + 3, '=');
  }

  if (count == 2) {
    wapp_str8_set(out, iout, char_at(buf[0] >> 2));
    wapp_str8_set(out, iout + 1, char_at(((buf[0] & 0x03) << 4) + (buf[1] >> 4)));
    wapp_str8_set(out, iout + 2, char_at((buf[1] & 0x0f) << 2));
    wapp_str8_set(out, iout + 3, '=');

    iout += 4;
  }

  return out;
}

Str8 *decode(Allocator *allocator, Str8RO *input) {
  if (input->size == 0) {
    return NULL;
  }

  u64 n_out = calc_decode_length(input);
  Str8 *out = wapp_str8_alloc_and_fill_buf(allocator, n_out);
  if (!out) return NULL;

  u8 buf[4] = {0};
  u8 count  = 0;
  u64 iout  = 0;

  for (u64 i = 0; i < input->size; ++i) {
    buf[count] = char_index(wapp_str8_get(input, i));
    count += 1;
    if (count == 4) {
      wapp_str8_set(out, iout, (buf[0] << 2) + (buf[1] >> 4));
      if (buf[2] != 64) {
        wapp_str8_set(out, iout + 1, (buf[1] << 4) + (buf[2] >> 2));
      }
      if (buf[3] != 64) {
        wapp_str8_set(out, iout + 2, (buf[2] << 6) + buf[3]);
      }

      iout += 3;
      count = 0;
    }
  }

  return out;
}

internal c8 char_at(u8 index) {
  return wapp_str8_get(&alphabet, index);
}

internal u8 char_index(c8 c) {
  if (c == '=') {
      return 64;
  }

  u8 index = 0;
  for (u8 i = 0; i < 63; ++i) {
    if (char_at(i) == c) {
        break;
    }
    index += 1;
  }

  return index;
}

internal u64 calc_encode_length(Str8RO *input) {
  if (input->size < 3) {
    return 4;
  }

  u64 n_groups = (u64)(ceilf((f32)input->size / 3.0f));

  return n_groups * 4;
}

internal u64 calc_decode_length(Str8RO *input) {
  if (input->size < 4) {
    return 3;
  }

  u64 n_groups = (u64)(floorf((f32)input->size / 4.0f));

  u64 groups = n_groups * 3;

  for (u64 i = input->size; i > 0; --i) {
    if (wapp_str8_get(input, i) == '=') {
      groups -= 1;
    }
  }

  return groups;
}