#include "ehs/Base64.h"

namespace ehs
{
	Str_8 Base64::EncodeURL(const Str_8 &input)
	{
		UInt_64 input_length = input.Size();

		// Calculate the output length
		UInt_64 output_length = (input_length * 4 + 2) / 3;

		// Allocate memory for the output
		Str_8 result(output_length);

		// Loop through the input and fill the output
		for (int i = 0, j = 0; i < input_length;) {

			// Take first byte and shift right by 2 bits
			UInt_32 octet_a = (UInt_8)input[i++];
			UInt_32 octet_b = i < input_length ? (UInt_8)input[i++] : 0;
			UInt_32 octet_c = i < input_length ? (UInt_8)input[i++] : 0;

			UInt_32 triple = (octet_a << 16) + (octet_b << 8) + octet_c;

			// Encode the 24-bits into four 6-bits integers
			result[j++] = asciiUrl[(triple >> 3 * 6) & 0x3F];
			result[j++] = asciiUrl[(triple >> 2 * 6) & 0x3F];
			result[j++] = asciiUrl[(triple >> 1 * 6) & 0x3F];
			result[j++] = asciiUrl[(triple >> 0 * 6) & 0x3F];
		}

		return result;
	}

	Str_8 Base64::Encode(const Str_8 &input)
	{
		UInt_64 input_length = input.Size();

		// Calculate the output length
		UInt_64 output_length = 4 * ((input_length + 2) / 3);

		// Allocate memory for the output
		Str_8 result(output_length);

		// Loop through the input and fill the output
		for (int i = 0, j = 0; i < input_length;) {

			// Take first byte and shift right by 2 bits
			UInt_32 octet_a = i < input_length ? (UInt_8)input[i++] : 0;
			UInt_32 octet_b = i < input_length ? (UInt_8)input[i++] : 0;
			UInt_32 octet_c = i < input_length ? (UInt_8)input[i++] : 0;

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

			// Encode the 24-bits into four 6-bits integers
			result[j++] = ascii[(triple >> 3 * 6) & 0x3F];
			result[j++] = ascii[(triple >> 2 * 6) & 0x3F];
			result[j++] = ascii[(triple >> 1 * 6) & 0x3F];
			result[j++] = ascii[(triple >> 0 * 6) & 0x3F];
		}

		// Add padding '='
		if (input_length % 3 == 1) {
			result[output_length - 1] = '=';
			result[output_length - 2] = '=';
		} else if (input_length % 3 == 2) {
			result[output_length - 1] = '=';
		}

		return result;
	}

	Str_8 Base64::DecodeURL(const Str_8 &input)
	{
		Str_8 result(input.Size() * 3 / 4);

		UInt_64 remaining = input.Size();
		UInt_64 offsetIn = 0;
		UInt_64 offsetOut = 0;
		UInt_8 quartet[4];

		while (remaining)
		{
			if (remaining >= 4)
			{
				for (UInt_8 i = 0; i < 4; ++i)
				{
					if (!IsBase64URL(input[offsetIn + i]))
						return {};

					quartet[i] = FindURL(input[offsetIn + i]);
					if (quartet[i] == EHS_UINT_8_MAX)
						return {};
				}

				result[offsetOut++] = (Char_8)((((UInt_8*)&quartet)[0] << 2) | (((UInt_8*)&quartet)[1] >> 4));
				result[offsetOut++] = (Char_8)((((UInt_8*)&quartet)[1] << 4) | (((UInt_8*)&quartet)[2] >> 2));
				result[offsetOut++] = (Char_8)((((UInt_8*)&quartet)[2] << 6) |  ((UInt_8*)&quartet)[3]);

				offsetIn += 4;
				remaining -= 4;
			}
			else
			{
				for (UInt_8 i = 0; i < 4; ++i)
				{
					if (i < remaining)
					{
						if (!IsBase64URL(input[offsetIn + i]))
							return {};

						quartet[i] = FindURL(input[offsetIn + i]);
						if (quartet[i] == EHS_UINT_8_MAX)
							return {};
					}
					else
						quartet[i] = 0;
				}

				result[offsetOut++] = (Char_8)((quartet[0] << 2) | (quartet[1] >> 4));

				if (remaining == 3)
					result[offsetOut++] = (Char_8)((quartet[1] << 4) | (quartet[2] >> 2));

				offsetIn += remaining;
				remaining = 0;
			}
		}

		return result;
	}

	Str_8 Base64::Decode(const Str_8 &input)
	{
		Str_8 result(input.Size() * 3 / 4);

		UInt_64 remaining = input.Size();
		UInt_64 offsetIn = 0;
		UInt_64 offsetOut = 0;
		UInt_8 quartet[4];

		while (remaining)
		{
			if (remaining >= 4)
			{
				for (UInt_8 i = 0; i < 4; ++i)
				{
					if (!IsBase64(input[offsetIn + i]))
						return {};

					quartet[i] = Find(input[offsetIn + i]);
					if (quartet[i] == EHS_UINT_8_MAX)
						return {};
				}

				result[offsetOut++] = (Char_8)((((UInt_8*)&quartet)[0] << 2) | (((UInt_8*)&quartet)[1] >> 4));
				result[offsetOut++] = (Char_8)((((UInt_8*)&quartet)[1] << 4) | (((UInt_8*)&quartet)[2] >> 2));
				result[offsetOut++] = (Char_8)((((UInt_8*)&quartet)[2] << 6) |  ((UInt_8*)&quartet)[3]);

				offsetIn += 4;
				remaining -= 4;
			}
			else
			{
				for (UInt_8 i = 0; i < 4; ++i)
				{
					if (i < remaining)
					{
						if (!IsBase64(input[offsetIn + i]))
							return {};

						quartet[i] = Find(input[offsetIn + i]);
						if (quartet[i] == EHS_UINT_8_MAX)
							return {};
					}
					else
						quartet[i] = 0;
				}

				result[offsetOut++] = (Char_8)((quartet[0] << 2) | (quartet[1] >> 4));

				if (remaining == 3)
					result[offsetOut++] = (Char_8)((quartet[1] << 4) | (quartet[2] >> 2));

				offsetIn += remaining;
				remaining = 0;
			}
		}

		return result;
	}

	UInt_8 Base64::FindURL(const UInt_8 &c)
	{
		for (UInt_8 i = 0; i < (UInt_8)sizeof(asciiUrl) - 1; ++i)
			if (asciiUrl[i] == c)
				return i;

		return EHS_UINT_8_MAX;;
	}

	UInt_8 Base64::Find(const UInt_8 &c)
	{
		for (UInt_8 i = 0; i < (UInt_8)sizeof(ascii) - 1; ++i)
			if (ascii[i] == c)
				return i;
		
		return EHS_SINT_8_MAX;
	}

	bool Base64::IsBase64URL(const UInt_8 &c)
	{
		return (c > 47 && c < 58) || (c > 64 && c < 91) || (c > 96 && c < 123) || c == '-' || c == '_';
	}

	bool Base64::IsBase64(const UInt_8 &c)
	{
		return (c > 47 && c < 58) || (c > 64 && c < 91) || (c > 96 && c < 123) || c == '+' || c == '/';
	}
}