#include "base64.h"

// Thanks, ryyst https://stackoverflow.com/questions/342409/how-do-i-base64-encode-decode-in-c

#include <stdint.h>
#include <stdlib.h>
#include <string.h>

/*
 * Data
 */
static const int mod_table[] = {0, 2, 1};
static const char *encoding_table =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
static char decoding_table[256] = {};

/*
 * Private API
 */
static void build_decoding_table() {
    for (int i = 0; i < 64; ++i) {
        decoding_table[(uint8_t)encoding_table[i]] = i;
    }
}

/*
 * Public API
 */
char *base64_encode(const void *_data, uint32_t size) {
    uint32_t out_size = 4 * ((size + 2) / 3) + 1;
    char *out_data = malloc(out_size);
    if (!out_size) {
        return NULL;
    }
    const char *data = _data;
    for (int i = 0, j = 0; i < size;) {
        uint32_t octet_a = i < size ? (uint8_t)data[i++] : 0;
        uint32_t octet_b = i < size ? (uint8_t)data[i++] : 0;
        uint32_t octet_c = i < size ? (uint8_t)data[i++] : 0;

        uint32_t triple = (octet_a << 0x10) + (octet_b << 0x08) + octet_c;

        out_data[j++] = encoding_table[(triple >> 3 * 6) & 0x3F];
        out_data[j++] = encoding_table[(triple >> 2 * 6) & 0x3F];
        out_data[j++] = encoding_table[(triple >> 1 * 6) & 0x3F];
        out_data[j++] = encoding_table[(triple >> 0 * 6) & 0x3F];
    }
    for (int i = 0; i < mod_table[size % 3]; i++)
        out_data[out_size - 2 - i] = '=';
    out_data[out_size - 1] = '\0';
    return out_data;
}

void *base64_decode(const char *data, uint32_t *out_size) {
    if (decoding_table[0] == 0)
        build_decoding_table();
    uint32_t size = strlen(data);
    if (size % 4 != 0) return NULL;

    *out_size = size / 4 * 3;
    if (data[size - 1] == '=') (*out_size)--;
    if (data[size - 2] == '=') (*out_size)--;

    uint8_t *decoded_data = malloc(*out_size);
    if (decoded_data == NULL) return NULL;

    for (int i = 0, j = 0; i < size;) {
        uint32_t sextet_a = data[i] == '=' ? 0 & i++ : decoding_table[data[i++]];
        uint32_t sextet_b = data[i] == '=' ? 0 & i++ : decoding_table[data[i++]];
        uint32_t sextet_c = data[i] == '=' ? 0 & i++ : decoding_table[data[i++]];
        uint32_t sextet_d = data[i] == '=' ? 0 & i++ : decoding_table[data[i++]];

        uint32_t triple = (sextet_a << 3 * 6)
            + (sextet_b << 2 * 6)
            + (sextet_c << 1 * 6)
            + (sextet_d << 0 * 6);

        if (j < *out_size) decoded_data[j++] = (triple >> 2 * 8) & 0xFF;
        if (j < *out_size) decoded_data[j++] = (triple >> 1 * 8) & 0xFF;
        if (j < *out_size) decoded_data[j++] = (triple >> 0 * 8) & 0xFF;
    }
    return decoded_data;
}