/*
 * Copyright 2011-2016, Axel Dörfler, axeld@pinc-software.de.
 * Distributed under the terms of the MIT License.
 */


#include "Response.h"

#include <algorithm>
#include <stdlib.h>

#include <UnicodeChar.h>


#define TRACE_IMAP
#ifdef TRACE_IMAP
#	define TRACE(...) printf(__VA_ARGS__)
#else
#	define TRACE(...) ;
#endif


namespace IMAP {


// Note, the following alphabet is a modified base64; the '/' is replaced by
// a ',' here.
static const char kBase64Alphabet[64] = {
  'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O',
  'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
  'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
  'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
  '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
  '+', ','
};
static char kInverseBase64Alphabet[128];
static bool kInverseBase64Initialized = false;


RFC3501Encoding::RFC3501Encoding()
{
	if (!kInverseBase64Initialized) {
		// This is not thread safe, but it's not harmful
		for (size_t i = 0; i < sizeof(kBase64Alphabet); i++)
			kInverseBase64Alphabet[(int)kBase64Alphabet[i]] = i + 1;
		kInverseBase64Initialized = true;
	}
}


RFC3501Encoding::~RFC3501Encoding()
{
}


BString
RFC3501Encoding::Encode(const BString& clearText) const
{
	const char* clear = clearText.String();
	bool shifted = false;
	int32 bitsToWrite = 0;
	int32 sextet = 0;
	BString buffer;

	while (true) {
		uint32 c = BUnicodeChar::FromUTF8(&clear);
		if (c == 0)
			break;

		if (!shifted && c == '&')
			buffer += "&-";
		else if (c >= 0x20 && c <= 0x7e) {
			_Unshift(buffer, bitsToWrite, sextet, shifted);
			buffer += c;
		} else {
			// Enter shifted mode, encode in base64
			if (!shifted) {
				buffer += '&';
				shifted = true;
			}

			bitsToWrite += 16;
			while (bitsToWrite >= 6) {
				bitsToWrite -= 6;
				buffer += kBase64Alphabet[(sextet + (c >> bitsToWrite)) & 0x3f];
				sextet = 0;
			}
			sextet = (c << (6 - bitsToWrite)) & 0x3f;
		}
	}

	_Unshift(buffer, bitsToWrite, sextet, shifted);
	return buffer;
}


BString
RFC3501Encoding::Decode(const BString& encodedText) const
{
	int32 end = encodedText.Length();
	BString buffer;
	for (int32 i = 0; i < end; i++) {
		uint8 c = (uint8)encodedText.ByteAt(i);
		if (c == '&') {
			if (i < end - 1 && encodedText.ByteAt(i + 1) == '-') {
				// just add an ampersand
				buffer += '&';
				i++;
			} else {
				// base64 encoded chunk
				uint32 value = 0;
				int32 bitsRead = 0;
				while (true) {
					if (++i >= end)
						throw ParseException("Malformed base64!");

					c = encodedText.ByteAt(i);
					if (c == '-') {
						if (value != 0 || bitsRead >= 6)
							throw ParseException("Base64 encoding ends early!");
						break;
					}
					if (c >= 128)
						throw ParseException("Malformed base64!");
					int32 sextet = kInverseBase64Alphabet[c] - 1;
					if (sextet >= 0) {
						bitsRead += 6;
						if (bitsRead < 16) {
							value += sextet << (16 - bitsRead);
						} else {
							bitsRead -= 16;
							value += sextet >> bitsRead;
							_ToUTF8(buffer, value);

							// Move on to next character
							value = (sextet << (16 - bitsRead)) & 0xffff;
						}
					} else {
						buffer += c;
						if (value != 0 || bitsRead >= 6)
							throw ParseException("Malformed base64!");
						break;
					}
				}
			}
		} else
			buffer += c;
	}
	return buffer;
}


void
RFC3501Encoding::_ToUTF8(BString& string, uint32 c) const
{
	if (c < 0x80)
		string += (char)c;
	else if (c < 0x800) {
		string += 0xc0 | (c >> 6);
		string += 0x80 | (c & 0x3f);
	} else if (c < 0x10000) {
		string += 0xe0 | (c >> 12);
		string += 0x80 | ((c >> 6) & 0x3f);
		string += 0x80 | (c & 0x3f);
	} else if (c <= 0x10ffff) {
		string += 0xf0 | (c >> 18);
		string += 0x80 | ((c >> 12) & 0x3f);
		string += 0x80 | ((c >> 6) & 0x3f);
		string += 0x80 | (c & 0x3f);
	}
}


//!	Exit base64, or "shifted" mode.
void
RFC3501Encoding::_Unshift(BString& buffer, int32& bitsToWrite, int32& sextet,
	bool& shifted) const
{
	if (!shifted)
		return;

	if (bitsToWrite != 0)
		buffer += kBase64Alphabet[sextet];
	buffer += '-';
	sextet = 0;
	bitsToWrite = 0;
	shifted = false;
}


// #pragma mark -


ArgumentList::ArgumentList()
	:
	BObjectList<Argument>(5, true)
{
}


ArgumentList::~ArgumentList()
{
}


bool
ArgumentList::Contains(const char* string) const
{
	for (int32 i = 0; i < CountItems(); i++) {
		if (StringArgument* argument
				= dynamic_cast<StringArgument*>(ItemAt(i))) {
			if (argument->String().ICompare(string) == 0)
				return true;
		}
	}
	return false;
}


BString
ArgumentList::StringAt(int32 index) const
{
	if (index >= 0 && index < CountItems()) {
		if (StringArgument* argument
				= dynamic_cast<StringArgument*>(ItemAt(index)))
			return argument->String();
	}
	return "";
}


bool
ArgumentList::IsStringAt(int32 index) const
{
	if (index >= 0 && index < CountItems()) {
		if (dynamic_cast<StringArgument*>(ItemAt(index)) != NULL)
			return true;
	}
	return false;
}


bool
ArgumentList::EqualsAt(int32 index, const char* string) const
{
	return StringAt(index).ICompare(string) == 0;
}


ArgumentList&
ArgumentList::ListAt(int32 index) const
{
	if (index >= 0 && index < CountItems()) {
		if (ListArgument* argument = dynamic_cast<ListArgument*>(ItemAt(index)))
			return argument->List();
	}

	static ArgumentList empty;
	return empty;
}


bool
ArgumentList::IsListAt(int32 index) const
{
	return index >= 0 && index < CountItems()
		&& dynamic_cast<ListArgument*>(ItemAt(index)) != NULL;
}


bool
ArgumentList::IsListAt(int32 index, char kind) const
{
	if (index >= 0 && index < CountItems()) {
		if (ListArgument* argument = dynamic_cast<ListArgument*>(ItemAt(index)))
			return argument->Kind() == kind;
	}
	return false;
}


uint32
ArgumentList::NumberAt(int32 index) const
{
	return atoul(StringAt(index).String());
}


bool
ArgumentList::IsNumberAt(int32 index) const
{
	BString string = StringAt(index);
	for (int32 i = 0; i < string.Length(); i++) {
		if (!isdigit(string.ByteAt(i)))
			return false;
	}
	return string.Length() > 0;
}


BString
ArgumentList::ToString() const
{
	BString string;

	for (int32 i = 0; i < CountItems(); i++) {
		if (i > 0)
			string += ", ";
		string += ItemAt(i)->ToString();
	}
	return string;
}


// #pragma mark -


Argument::Argument()
{
}


Argument::~Argument()
{
}


// #pragma mark -


ListArgument::ListArgument(char kind)
	:
	fKind(kind)
{
}


BString
ListArgument::ToString() const
{
	BString string(fKind == '[' ? "[" : "(");
	string += fList.ToString();
	string += fKind == '[' ? "]" : ")";

	return string;
}


// #pragma mark -


StringArgument::StringArgument(const BString& string)
	:
	fString(string)
{
}


StringArgument::StringArgument(const StringArgument& other)
	:
	fString(other.fString)
{
}


BString
StringArgument::ToString() const
{
	return fString;
}


// #pragma mark -


ParseException::ParseException()
{
	fBuffer[0] = '\0';
}


ParseException::ParseException(const char* format, ...)
{
	va_list args;
	va_start(args, format);
	vsnprintf(fBuffer, sizeof(fBuffer), format, args);
	va_end(args);
}


// #pragma mark -


StreamException::StreamException(status_t status)
	:
	fStatus(status)
{
}


// #pragma mark -


ExpectedParseException::ExpectedParseException(char expected, char instead)
{
	char bufferA[8];
	char bufferB[8];
	snprintf(fBuffer, sizeof(fBuffer), "Expected %s, but got %s instead!",
		CharToString(bufferA, sizeof(bufferA), expected),
		CharToString(bufferB, sizeof(bufferB), instead));
}


const char*
ExpectedParseException::CharToString(char* buffer, size_t size, char c)
{
	snprintf(buffer, size, isprint(c) ? "\"%c\"" : "(%x)", c);
	return buffer;
}


// #pragma mark -


LiteralHandler::LiteralHandler()
{
}


LiteralHandler::~LiteralHandler()
{
}


// #pragma mark -


Response::Response()
	:
	fTag(0),
	fContinuation(false),
	fHasNextChar(false)
{
}


Response::~Response()
{
}


void
Response::Parse(BDataIO& stream, LiteralHandler* handler)
{
	MakeEmpty();
	fLiteralHandler = handler;
	fTag = 0;
	fContinuation = false;
	fHasNextChar = false;

	char begin = Next(stream);
	if (begin == '*') {
		// Untagged response
		Consume(stream, ' ');
	} else if (begin == '+') {
		// Continuation
		fContinuation = true;
	} else if (begin == 'A') {
		// Tagged response
		fTag = ExtractNumber(stream);
		Consume(stream, ' ');
	} else
		throw ParseException("Unexpected response begin");

	char c = ParseLine(*this, stream);
	if (c != '\0')
		throw ExpectedParseException('\0', c);
}


bool
Response::IsCommand(const char* command) const
{
	return IsUntagged() && EqualsAt(0, command);
}


char
Response::ParseLine(ArgumentList& arguments, BDataIO& stream)
{
	while (true) {
		char c = Peek(stream);
		if (c == '\0')
			break;

		switch (c) {
			case '(':
				ParseList(arguments, stream, '(', ')');
				break;
			case '[':
				ParseList(arguments, stream, '[', ']');
				break;
			case ')':
			case ']':
				Consume(stream, c);
				return c;
			case '"':
				ParseQuoted(arguments, stream);
				break;
			case '{':
				ParseLiteral(arguments, stream);
				break;

			case ' ':
			case '\t':
				// whitespace
				Consume(stream, c);
				break;

			case '\r':
				Consume(stream, '\r');
				Consume(stream, '\n');
				return '\0';
			case '\n':
				Consume(stream, '\n');
				return '\0';

			default:
				ParseString(arguments, stream);
				break;
		}
	}

	return '\0';
}


void
Response::ParseList(ArgumentList& arguments, BDataIO& stream, char start,
	char end)
{
	Consume(stream, start);

	ListArgument* argument = new ListArgument(start);
	arguments.AddItem(argument);

	char c = ParseLine(argument->List(), stream);
	if (c != end)
		throw ExpectedParseException(end, c);
}


void
Response::ParseQuoted(ArgumentList& arguments, BDataIO& stream)
{
	Consume(stream, '"');

	BString string;
	while (true) {
		char c = Next(stream);
		if (c == '\\') {
			c = Next(stream);
		} else if (c == '"') {
			arguments.AddItem(new StringArgument(string));
			return;
		}
		if (c == '\0')
			break;

		string += c;
	}

	throw ParseException("Unexpected end of qouted string!");
}


void
Response::ParseLiteral(ArgumentList& arguments, BDataIO& stream)
{
	Consume(stream, '{');
	size_t size = ExtractNumber(stream);
	Consume(stream, '}');
	Consume(stream, '\r');
	Consume(stream, '\n');

	bool handled = false;
	if (fLiteralHandler != NULL) {
		handled = fLiteralHandler->HandleLiteral(*this, arguments, stream,
			size);
	}

	if (!handled && size <= 65536) {
		// The default implementation just adds the data as a string
		TRACE("Trying to read literal with %" B_PRIuSIZE " bytes.\n", size);
		BString string;
		char* buffer = string.LockBuffer(size);
		if (buffer == NULL) {
			throw ParseException("Not enough memory for literal of %"
				B_PRIuSIZE " bytes.", size);
		}

		size_t totalRead = 0;
		while (totalRead < size) {
			ssize_t bytesRead = stream.Read(buffer + totalRead,
				size - totalRead);
			if (bytesRead == 0)
				throw StreamException(B_IO_ERROR);
			if (bytesRead < 0)
				throw StreamException(bytesRead);

			totalRead += bytesRead;
		}

		string.UnlockBuffer(size);
		arguments.AddItem(new StringArgument(string));
	} else {
		// Skip any bytes left in the literal stream
		_SkipLiteral(stream, size);
	}
}


void
Response::ParseString(ArgumentList& arguments, BDataIO& stream)
{
	arguments.AddItem(new StringArgument(ExtractString(stream)));
}


BString
Response::ExtractString(BDataIO& stream)
{
	BString string;

	// TODO: parse modified UTF-7 as described in RFC 3501, 5.1.3
	while (true) {
		char c = Peek(stream);
		if (c == '\0')
			break;
		if (c <= ' ' || strchr("()[]{}\"", c) != NULL)
			return string;

		string += Next(stream);
	}

	throw ParseException("Unexpected end of string");
}


size_t
Response::ExtractNumber(BDataIO& stream)
{
	BString string = ExtractString(stream);

	const char* end;
	size_t number = strtoul(string.String(), (char**)&end, 10);
	if (end == NULL || end[0] != '\0')
		throw ParseException("Invalid number!");

	return number;
}


void
Response::Consume(BDataIO& stream, char expected)
{
	char c = Next(stream);
	if (c != expected)
		throw ExpectedParseException(expected, c);
}


char
Response::Next(BDataIO& stream)
{
	if (fHasNextChar) {
		fHasNextChar = false;
		return fNextChar;
	}

	return Read(stream);
}


char
Response::Peek(BDataIO& stream)
{
	if (fHasNextChar)
		return fNextChar;

	fNextChar = Read(stream);
	fHasNextChar = true;

	return fNextChar;
}


char
Response::Read(BDataIO& stream)
{
	char c;
	ssize_t bytesRead = stream.Read(&c, 1);
	if (bytesRead == 1) {
		printf("%c", c);
		return c;
	}

	if (bytesRead == 0)
		throw StreamException(B_IO_ERROR);

	throw StreamException(bytesRead);
}


void
Response::_SkipLiteral(BDataIO& stream, size_t size)
{
	char buffer[4096];
	size_t totalRead = 0;
	while (totalRead < size) {
		size_t toRead = std::min(sizeof(buffer), size - totalRead);
		ssize_t bytesRead = stream.Read(buffer, toRead);
		if (bytesRead == 0)
			throw StreamException(B_IO_ERROR);
		if (bytesRead < 0)
			throw StreamException(bytesRead);

		totalRead += bytesRead;
	}
}


// #pragma mark -


ResponseParser::ResponseParser(BDataIO& stream)
	:
	fLiteralHandler(NULL)
{
	SetTo(stream);
}


ResponseParser::~ResponseParser()
{
}


void
ResponseParser::SetTo(BDataIO& stream)
{
	fStream = &stream;
}


void
ResponseParser::SetLiteralHandler(LiteralHandler* handler)
{
	fLiteralHandler = handler;
}


status_t
ResponseParser::NextResponse(Response& response, bigtime_t timeout)
{
	response.Parse(*fStream, fLiteralHandler);
	return B_OK;
}


}	// namespace IMAP