#include "base64.h"

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

static const char* alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

char* base64_strict_encode( void* v_binary, size_t len )
{
	// RFC 4648 - https://www.rfc-editor.org/rfc/rfc4648.html
	enum { pad = '=' };

	uint8_t* binary = v_binary;
	const int groups = ( len + 2 ) / 3;

	char* result = malloc( groups * 4 + 1 );
	result[groups*4] = '\0';

	uint32_t word;
	char* out = result;
	for( int remain = len; remain > 0; remain -= 3, binary += 3, out += 4 ) {
		word = 0;
		switch( remain ) {
			default:
				word |= binary[2] << (8*0);
			case 2:
				word |= binary[1] << (8*1);
			case 1:
				word |= binary[0] << (8*2);
		}
		switch( remain ) {
			default:
				out[3] = alphabet[ ( word >> (6*0) ) % 64 ];
			case 2:
				out[2] = alphabet[ ( word >> (6*1) ) % 64 ];
			case 1:
				out[1] = alphabet[ ( word >> (6*2) ) % 64 ];
				out[0] = alphabet[ ( word >> (6*3) ) % 64 ];
		}
		switch( remain ) {
			case 1:
				out[3] = '=';
				out[2] = '=';
				break;
			case 2:
				out[3] = '=';
				break;
		}
	}

	return result;
}
bool base64_decode( const char* str, void** v_binary, size_t* len )
{
	int input_length = strlen(str);

	uint8_t* data = malloc( ( input_length * 6 + 7 ) / 8 );
	*v_binary = data;
	*len = 0;

	uint32_t buffer = 0;
	int bits = 0;
	uint8_t part;

	for(;;) {
		//printf( "  - Input %02X (%c)\n", *str, *str );
		switch( *str ) {
			case 'A' ... 'Z': part = *str - 'A'; break;
			case 'a' ... 'z': part = *str - 'a' + 26; break;
			case '0' ... '9': part = *str - '0' + 52; break;
			case '+':         part = 62; break;
			case '/':         part = 63; break;
			case '=':         goto finish;
			case '\0':
				goto done;
			default:
				free( *v_binary );
				*len = 0;
				return false;
		};
		++str;

		buffer = ( buffer << 6 ) | part;
		bits += 6;
		//printf( "buffer=%08X, bits=%d\n", buffer, bits );

		if( bits > 8 ) {
			uint8_t byte = buffer >> ( bits % 8 );
			*data = byte;
			//printf( "  - Output: %02X (%c)\n", byte, byte );
			buffer &= ( 1 << (bits%8) ) - 1;
			*len += 1;
			++data;
			bits -= 8;
		}
	}
done:
	*len += 1;
finish:
	uint8_t byte = buffer >> ( bits % 8 );
	*data = byte;
	++data;
	bits = 0;
	return true;
}