package main

import "core:math"
import "core:strings"

@(private="file")
alphabet := "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

encode :: proc(s: string) -> string {
  if len(s) == 0 do return "";

  n_out := calc_encode_length(s);

  bytes := make([]u8, n_out);
  defer delete(bytes);

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

  for i in 0..<len(s) {
    buf[count] = s[i];
    count += 1;
    if count == 3 {
      bytes[iout] = char_at(buf[0] >> 2);
      bytes[iout + 1] = char_at(((buf[0] & 0x03) << 4) + (buf[1] >> 4));
      bytes[iout + 2] = char_at(((buf[1] & 0x0f) << 2) + (buf[2] >> 6));
      bytes[iout + 3] = char_at(buf[2] & 0x3f);

      iout += 4;
      count = 0;
    }
  }

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

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

    iout += 4;
  }

  output, err := strings.clone_from_bytes(bytes[:]);
  if err != nil {
    return "";
  }

  return output;
}

decode :: proc(s: string) -> string {
  if len(s) == 0 do return "";

  n_out := calc_decode_length(s);

  bytes := make([]u8, n_out);
  defer delete(bytes);

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

  for i in 0..<len(s) {
    buf[count] = char_index(s[i]);
    count += 1;
    if count == 4 {
      bytes[iout] = (buf[0] << 2) + (buf[1] >> 4);
      if buf[2] != 64 {
        bytes[iout + 1] = (buf[1] << 4) + (buf[2] >> 2);
      }
      if buf[3] != 64 {
        bytes[iout + 2] = (buf[2] << 6) + buf[3];
      }

      iout += 3;
      count = 0;
    }
  }

  output, err := strings.clone_from_bytes(bytes[:]);
  if err != nil {
    return "";
  }

  return output;
}

@(private="file")
char_at :: proc(index: u8) -> u8 {
  return alphabet[index];
}

@(private="file")
char_index :: proc(char: u8) -> u8 {
  if char == '=' {
    return 64;
  }

  index : u8 = 0;
  for c in alphabet {
    if c == rune(char) do break;
    index += 1;
  }

  return index;
}

@(private="file")
calc_encode_length :: proc(s: string) -> u64 {
  if len(s) < 3 {
    return 4;
  }

  groups := u64(math.ceil_f32(f32(len(s)) / 3.0));

  return groups * 4;
}

@(private="file")
calc_decode_length :: proc(s: string) -> u64 {
  if len(s) < 4 {
    return 3;
  }

  n_groups := u64(math.floor_f32(f32(len(s)) / 4.0));
  groups := n_groups * 3;
  #reverse for char in s {
    if char == '=' {
      groups -= 1;
    }
  }

  return groups;
}