#include <libcper/base64.h>
#include <libcper/BaseTypes.h>

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <assert.h>

static const UINT8 encode_table[65] =
	"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

/**
 *
 * Caller is responsible for freeing the returned buffer.
 */
CHAR8 *base64_encode(const UINT8 *src, INT32 len, INT32 *out_len)
{
	CHAR8 *out;
	CHAR8 *out_pos;
	const UINT8 *src_end;
	const UINT8 *in_pos;

	if (len <= 0) {
		return NULL;
	}

	if (!out_len) {
		return NULL;
	}

	// 3 byte blocks to 4 byte blocks plus up to 2 bytes of padding
	*out_len = 4 * ((len + 2) / 3);

	// Handle overflows
	if (*out_len < len) {
		return NULL;
	}

	out = malloc(*out_len);
	if (out == NULL) {
		return NULL;
	}

	src_end = src + len;
	in_pos = src;
	out_pos = out;
	while (src_end - in_pos >= 3) {
		*out_pos++ = encode_table[in_pos[0] >> 2];
		*out_pos++ = encode_table[((in_pos[0] & 0x03) << 4) |
					  (in_pos[1] >> 4)];
		*out_pos++ = encode_table[((in_pos[1] & 0x0f) << 2) |
					  (in_pos[2] >> 6)];
		*out_pos++ = encode_table[in_pos[2] & 0x3f];
		in_pos += 3;
	}

	if (src_end - in_pos) {
		*out_pos++ = encode_table[in_pos[0] >> 2];
		if (src_end - in_pos == 1) {
			*out_pos++ = encode_table[(in_pos[0] & 0x03) << 4];
			*out_pos++ = '=';
		} else {
			*out_pos++ = encode_table[((in_pos[0] & 0x03) << 4) |
						  (in_pos[1] >> 4)];
			*out_pos++ = encode_table[(in_pos[1] & 0x0f) << 2];
		}
		*out_pos++ = '=';
	}

	return out;
}

// Base64 decode table.  Invalid values are specified with 0x80.
UINT8 decode_table[256] = {
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x3e, 0x80, 0x80, 0x80, 0x3f,
	0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a, 0x3b, 0x3c, 0x3d, 0x80, 0x80,
	0x80, 0x00, 0x80, 0x80, 0x80, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06,
	0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12,
	0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20, 0x21, 0x22, 0x23, 0x24,
	0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30,
	0x31, 0x32, 0x33, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80,
	0x80, 0x80, 0x80, 0x80,
};

/**
 *
 * Caller is responsible for freeing the returned buffer.
 */
UINT8 *base64_decode(const CHAR8 *src, INT32 len, INT32 *out_len)
{
	UINT8 *out = NULL;
	UINT8 *pos = NULL;
	UINT8 block[4];
	INT32 block_index = 0;
	INT32 src_index = 0;

	if (!out_len) {
		goto error;
	}

	// Malloc might be up to 2 larger dependent on padding
	*out_len = len / 4 * 3 + 2;
	pos = out = malloc(*out_len);
	if (out == NULL) {
		goto error;
	}

	block_index = 0;
	for (src_index = 0; src_index < len; src_index++) {
		char current_char = src[src_index];
		if (current_char == '=') {
			break;
		}
		// If the final character is a newline, as can occur in many editors
		// then ignore it.
		if (src_index + 1 == len && current_char == '\n') {
			printf("Ignoring trailing newline.\n");
			break;
		}

		block[block_index] = decode_table[(UINT8)current_char];
		if (block[block_index] == 0x80) {
			printf("Invalid character \"%c\".\n", current_char);
			goto error;
		}

		block_index++;
		if (block_index == 4) {
			*pos++ = (block[0] << 2) | (block[1] >> 4);
			*pos++ = (block[1] << 4) | (block[2] >> 2);
			*pos++ = (block[2] << 6) | block[3];
			block_index = 0;
		}
	}
	if (block_index == 0) {
		// mod 4 Even number of characters, no padding.
	} else if (block_index == 1) {
		printf("Invalid base64 input length.  Last character truncated.\n");
		goto error;
	} else if (block_index == 2) {
		*pos++ = (block[0] << 2) | (block[1] >> 4);
	} else if (block_index == 3) {
		*pos++ = (block[0] << 2) | (block[1] >> 4);
		*pos++ = (block[1] << 4) | (block[2] >> 2);
	} else {
		/* Invalid pad_counting */
		printf("Invalid base64 input length %d.\n", block_index);
		goto error;
	}

	*out_len = pos - out;
	return out;

error:
	free(out);
	return NULL;
}