#include "Base64.hh"
#include "xrange.hh"
#include <algorithm>
#include <cassert>
#include <cstdint>

namespace Base64 {

using std::string;

static inline char encode(uint8_t c)
{
	static const char* const base64_chars =
		"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
		"abcdefghijklmnopqrstuvwxyz"
		"0123456789+/";
	assert(c < 64);
	return base64_chars[c];
}

static inline uint8_t decode(unsigned char c)
{
	if        ('A' <= c && c <= 'Z') {
		return c - 'A';
	} else if ('a' <= c && c <= 'z') {
		return c - 'a' + 26;
	} else if ('0' <= c && c <= '9') {
		return c - '0' + 52;
	} else if (c == '+') {
		return 62;
	} else if (c == '/') {
		return 63;
	} else {
		return uint8_t(-1);
	}
}

string encode(const void* input_, size_t inSize)
{
	static const int CHUNKS = 19;
	static const int IN_CHUNKS  = 3 * CHUNKS;
	static const int OUT_CHUNKS = 4 * CHUNKS; // 76 chars per line

	auto input = static_cast<const uint8_t*>(input_);
	auto outSize = ((inSize + (IN_CHUNKS - 1)) / IN_CHUNKS) * (OUT_CHUNKS + 1); // overestimation
	string ret(outSize, 0); // too big

	size_t out = 0;
	while (inSize) {
		if (out) ret[out++] = '\n';
		auto n2 = std::min<size_t>(IN_CHUNKS, inSize);
		auto n = unsigned(n2);
		for (/**/; n >= 3; n -= 3) {
			ret[out++] = encode( (input[0] & 0xfc) >> 2);
			ret[out++] = encode(((input[0] & 0x03) << 4) +
			                    ((input[1] & 0xf0) >> 4));
			ret[out++] = encode(((input[1] & 0x0f) << 2) +
			                    ((input[2] & 0xc0) >> 6));
			ret[out++] = encode( (input[2] & 0x3f) >> 0);
			input += 3;
		}
		if (n) {
			uint8_t buf3[3] = { 0, 0, 0 };
			for (unsigned i = 0; i < n; ++i) {
				buf3[i] = input[i];
			}
			uint8_t buf4[4];
			buf4[0] =  (buf3[0] & 0xfc) >> 2;
			buf4[1] = ((buf3[0] & 0x03) << 4) +
				  ((buf3[1] & 0xf0) >> 4);
			buf4[2] = ((buf3[1] & 0x0f) << 2) +
				  ((buf3[2] & 0xc0) >> 6);
			buf4[3] =  (buf3[2] & 0x3f) >> 0;
			for (unsigned j = 0; (j < n + 1); ++j) {
				ret[out++] = encode(buf4[j]);
			}
			for (/**/; n < 3; ++n) {
				ret[out++] = '=';
			}
		}
		inSize -= n2;
	}

	assert(outSize >= out);
	ret.resize(out); // shrink to correct size
	return ret;
}

string decode(const string& input)
{
	auto inSize = input.size();
	auto outSize = (inSize * 3 + 3) / 4; // overestimation
	string ret(outSize, 0); // too big

	unsigned i = 0;
	size_t out = 0;
	uint8_t buf4[4];
	for (auto in : xrange(inSize)) {
		uint8_t d = decode(input[in]);
		if (d == uint8_t(-1)) continue;
		buf4[i++] = d;
		if (i == 4) {
			i = 0;
			ret[out++] = char(((buf4[0] & 0xff) << 2) + ((buf4[1] & 0x30) >> 4));
			ret[out++] = char(((buf4[1] & 0x0f) << 4) + ((buf4[2] & 0x3c) >> 2));
			ret[out++] = char(((buf4[2] & 0x03) << 6) + ((buf4[3] & 0xff) >> 0));
		}
	}
	if (i) {
		for (unsigned j = i; j < 4; ++j) {
			buf4[j] = 0;
		}
		uint8_t buf3[3];
		buf3[0] = ((buf4[0] & 0xff) << 2) + ((buf4[1] & 0x30) >> 4);
		buf3[1] = ((buf4[1] & 0x0f) << 4) + ((buf4[2] & 0x3c) >> 2);
		buf3[2] = ((buf4[2] & 0x03) << 6) + ((buf4[3] & 0xff) >> 0);
		for (unsigned j = 0; (j < i - 1); ++j) {
			ret[out++] = buf3[j];
		}
	}

	assert(outSize >= out);
	ret.resize(out); // shrink to correct size
	return ret;
}

} // namespace Base64
