#pragma once

#include <string>
#include <vector>
#include <cmath>
#include "flow/flow.h"
#include "flow/Trace.h"
#include "fdbclient/JSONDoc.h"

class JsonBuilder;
class JsonBuilderObject;
class JsonBuilderArray;
typedef JsonBuilder JsonString;
template <typename T> class JsonBuilderObjectSetter;

// Class for building JSON string values.
// Default value is null, as in the JSON type
class JsonBuilder {
protected:
	enum EType { NULLVALUE, OBJECT, ARRAY };

	typedef VectorRef<char> VString;
public:
	// Default value is null, which will be considered "empty"
	JsonBuilder() : type(NULLVALUE), elements(0), bytes(0) {
		jsonText.resize(arena, 1);
	}

	int getFinalLength() const {
		return bytes + strlen(getEnd());
	}

	// TODO: Remove the need for this by changing usages to steal this's content
	std::string getJson() const {
		std::string result;
		result.reserve(bytes + 1);
		for(auto& it : jsonText) {
			result.append(it.begin(), it.end());
		}
		result.append(getEnd());
		return result;
	}

	int size() const {
		return elements;
	}

	bool empty() const {
		return elements == 0;
	}

	static JsonBuilderObject makeMessage(const char *name, const char *description);

	static int coerceAsciiNumberToJSON(const char *s, int len, char *dst);

protected:
	EType type;
	Arena arena;
	mutable VectorRef<VString> jsonText;
	int elements;
	int bytes;

	// 'raw' write methods
	inline void write(const char *s, int len) {
		bytes += len;
		jsonText.back().append(arena, s, len);
	}

	inline void write(const char* s) {
		write(s, strlen(s));
	}

	inline void write(const StringRef &s) {
		write((char *)s.begin(), s.size());
	}

	inline void write(char s) {
		++bytes;
		jsonText.back().push_back(arena, s);
	}

	// writeValue() methods write JSON form of the value
	void writeValue(const json_spirit::mValue &val) {
		switch(val.type()) {
			case json_spirit::int_type:
				return writeValue(val.get_int64());
			case json_spirit::bool_type:
				return writeValue(val.get_bool());
			case json_spirit::real_type:
				return writeValue(val.get_real());
			case json_spirit::str_type:
				return writeValue(val.get_str());
			default:
				// Catch-all for objects/arrays
				return write(json_spirit::write_string(val));
		};
	}

	void writeValue(const bool& val) {
		write(val ? "true" : "false");
	}

	template<typename T> inline void writeFormat(const char *fmt, const T &val) {
		VString &dst = jsonText.back();
		const int limit = 30;
		dst.reserve(arena, dst.size() + limit);
		int len = snprintf(dst.end(), limit, fmt, val);
		if(len > 0 && len < limit) {
			dst.extendUnsafeNoReallocNoInit(len);
		}
		else {
			write(format(fmt, val));
		}
	}

	void writeValue(const int64_t& val) {
		writeFormat("%lld", val);
	}

	void writeValue(const uint64_t& val) {
		writeFormat("%llu", val);
	}

	void writeValue(const int& val) {
		writeFormat("%d", val);
	}

	void writeValue(const double& val) {
		if(std::isfinite(val)) {
			writeFormat("%g", val);
		}
		else if(std::isnan(val)) {
			write("-999");
		}
		else {
			write("1e99");
		}
	}

	bool shouldEscape(char c) {
		switch( c ) {
			case '"':
			case '\\':
			case '\b':
			case '\f':
			case '\n':
			case '\r':
			case '\t':
				return true;
			default:
				return false;
		}
	}

	void writeValue(const char *val, int len) {
		write('"');
		int beginCopy = 0;
		VString &dst = jsonText.back();
		for (int i = 0; i < len; i++) {
			if (shouldEscape(val[i])) {
				dst.append(arena, val + beginCopy, i - beginCopy);
				beginCopy = i + 1;
				write('\\');
				write(val[i]);
			}
		}
		if(beginCopy < len) {
			dst.append(arena, val + beginCopy, len - beginCopy);
		}
		write('"');
	}

	inline void writeValue(const std::string& val) {
		writeValue(val.data(), val.size());
	}

	inline void writeValue(const char* val) {
		writeValue(val, strlen(val));
	}

	inline void writeValue(const StringRef &s) {
		writeValue((const char *)s.begin(), s.size());
	}

	// Write the finalized (closed) form of val
	void writeValue(const JsonBuilder &val) {
		bytes += val.bytes;
		jsonText.append(arena, val.jsonText.begin(), val.jsonText.size());
		val.jsonText.push_back(arena, VString());
		arena.dependsOn(val.arena);
		write(val.getEnd());
	}

	void writeCoercedAsciiNumber(const char *s, int len) {
		VString &val = jsonText.back();
		val.reserve(arena, val.size() + len + 3);
		int written = coerceAsciiNumberToJSON(s, len, val.end());
		if(written > 0) {
			val.extendUnsafeNoReallocNoInit(written);
		}
		else {
			write("-999");
		}
	}

	inline void writeCoercedAsciiNumber(const StringRef &s) {
		writeCoercedAsciiNumber((const char *)s.begin(), s.size());
	}

	inline void writeCoercedAsciiNumber(const std::string &s) {
		writeCoercedAsciiNumber(s.data(), s.size());
	}

	// Helper function to add contents of another JsonBuilder to this one.
	// This is only used by the subclasses to combine like-typed (at compile time) objects,
	// so it can be assumed that the other object has been initialized with an opening character.
	void _addContents(const JsonBuilder &other) {
		if(other.empty()) {
			return;
		}

		if(elements > 0) {
			write(',');
		}

		// Add everything but the first byte of the first string in arr
		bytes += other.bytes - 1;
		const VString &front = other.jsonText.front();
		jsonText.push_back(arena, front.slice(1, front.size()));
		jsonText.append(arena, other.jsonText.begin() + 1, other.jsonText.size() - 1);

		// Both JsonBuilders would now want to write to the same additional VString capacity memory
		// if they were modified, so force the other (likely to not be modified again) to start a new one.
		other.jsonText.push_back(arena, VString());

		arena.dependsOn(other.arena);
		elements += other.elements;
	}

	// Get the text necessary to finish the JSON string
	const char * getEnd() const {
		switch(type) {
			case NULLVALUE:
				return "null";
			case OBJECT:
				return "}";
			case ARRAY:
				return "]";
			default:
				return "";
		};
	}
};

class JsonBuilderArray : public JsonBuilder {
public:
	JsonBuilderArray() {
		type = ARRAY;
		write('[');
	}

	template<typename VT> inline JsonBuilderArray & push_back(const VT &val) {
		if(elements++ > 0) {
			write(',');
		}
		writeValue(val);
		return *this;
	}

	JsonBuilderArray & addContents(const json_spirit::mArray &arr) {
		for(auto &v : arr) {
			push_back(v);
		}
		return *this;
	}

	JsonBuilderArray & addContents(const JsonBuilderArray &arr) {
		_addContents(arr);
		return *this;
	}
};

class JsonBuilderObject : public JsonBuilder {
public:
	JsonBuilderObject() {
		type = OBJECT;
		write('{');
	}

	template<typename KT, typename VT> inline JsonBuilderObject & setKey(const KT &name, const VT &val) {
		if(elements++ > 0) {
			write(',');
		}
		write('"');
		write(name);
		write("\":");
		writeValue(val);
		return *this;
	}

	template<typename KT, typename VT> inline JsonBuilderObject & setKeyRawNumber(const KT &name, const VT &val) {
		if(elements++ > 0) {
			write(',');
		}
		write('"');
		write(name);
		write("\":");
		writeCoercedAsciiNumber(val);
		return *this;
	}

	template<typename T> inline JsonBuilderObjectSetter<T> operator[](T &&name);

	JsonBuilderObject & addContents(const json_spirit::mObject &obj) {
		for(auto &kv : obj) {
			setKey(kv.first, kv.second);
		}
		return *this;
	}

	JsonBuilderObject & addContents(const JsonBuilderObject &obj) {
		_addContents(obj);
		return *this;
	}

};

// Template is the key name, accepted as an r-value if possible to avoid copying if it's a string
template<typename KT>
class JsonBuilderObjectSetter {
public:
	JsonBuilderObjectSetter(JsonBuilderObject &dest, KT &&name) : dest(dest), name(std::forward<KT>(name)) {}

	// Value is accepted as an rvalue if possible
	template <class VT> inline void operator=(const VT &value) {
		dest.setKey(name, value);
	}

protected:
	JsonBuilderObject& dest;
	KT name;
};

template<typename T> inline JsonBuilderObjectSetter<T> JsonBuilderObject::operator[](T &&name) {
	return JsonBuilderObjectSetter<T>(*this, std::forward<T>(name));
}

