use lang:asm;

/**
 * Class that represents a JSON value.
 *
 * The value is essentially a union of the following types:
 * - Null
 * - Bools
 * - Number (int/float)
 * - String
 * - Array
 * - Map (string to value)
 *
 * Newly created instances are "null" by default. This value will happily be converted into any of
 * the supported types based on how it is used. Also, the number type is a bit flexible to match
 * usage in Storm. For example, numbers are kept as integers as long as possible, but will act as if
 * they are floats/doubles. Once the value is initialized to some type, it will throw an exception
 * if it is used as the incorrect type.
 *
 * Note: this class is explicitly designed to "go around" the null-safety of Storm, and instead
 * relies on exceptions. This is to make it easier to navigate deep data structures in a concise
 * manner.
 */
class JsonValue {
	// Default constructor, initialized to "null", which always throws when accessed until it is
	// initialized to some value.
	init() {
		init { type = Type:empty; }
	}

	// Initialize to a bool.
	cast(Bool v) {
		init {
			type = Type:bool;
			numberVal = if (v) { 1l; } else { 0l; };
		}
	}

	// Initialize to an integer.
	cast(Byte v) { self(v.long); }
	cast(Int v) { self(v.long); }
	cast(Nat v) { self(v.long); }
	cast(Word v) { self(v.long); }
	cast(Long v) {
		init {
			type = Type:long;
			numberVal = v;
		}
	}

	// Initialize to a floating point number.
	cast(Float v) { self(v.double); }
	cast(Double v) {
		Long converted;
		asm { mov converted, v; }
		init {
			type = Type:double;
			numberVal = converted;
		}
	}

	// Initialize to a string.
	cast(Str v) {
		init {
			type = Type:string;
			objectVal = v;
		}
	}

	// Initialize to an array.
	init(JsonValue[] v) {
		init {
			type = Type:array;
			objectVal = v;
		}
	}

	// Initialize to a map (object).
	init(Str->JsonValue v) {
		init {
			type = Type:object;
			objectVal = v;
		}
	}

	// Create initialized to an empty array.
	JsonValue emptyArray() : static {
		JsonValue(JsonValue[]);
	}

	// Create initialized to an empty object.
	JsonValue emptyObject() : static {
		JsonValue(Str->JsonValue);
	}

	// Compare to another value.
	Bool ==(JsonValue other) {
		if (type != other.type)
			return false;

		if (type == Type:bool) {
			return compare(numberVal != 0, other.numberVal != 0);
		} else if (type == Type:long) {
			return compare(numberVal, other.numberVal);
		} else if (type == Type:double) {
			// TODO: Maybe with some epsilon?
			return compare(numberVal, other.numberVal);
		} else if (type == Type:empty) {
			return true;
		} else if (objectVal as Str) {
			if (o = other.objectVal as Str)
				return compare(objectVal, o);
			else
				return false;
		} else if (objectVal as JsonValue[]) {
			if (o = other.objectVal as JsonValue[])
				return compare(objectVal, o);
			else
				return false;
		} else if (objectVal as Str->JsonValue) {
			if (o = other.objectVal as Str->JsonValue)
				return compare(objectVal, o);
			else
				return false;
		} else {
			// Inconsistent representation.
			return false;
		}
	}

	// Does this class contain a non-null value?
	Bool any() { type != Type:empty; }
	Bool empty() { type == Type:empty; }

	// Check type manually.
	Bool isNull() { type == Type:empty; }
	Bool isBool() { type == Type:bool; }
	Bool isNumber() { (type == Type:long) | (type == Type:double); }
	Bool isInteger() { type == Type:double; }
	Bool isStr() { type == Type:string; }
	Bool isArray() { type == Type:array; }
	Bool isObject() { type == Type:object; }

	// Type conversions. Throws if the type is incorrect.
	Bool bool() {
		if (type != Type:bool)
			throw JsonAccessError("Expected bool", this);
		numberVal != 0;
	}
	Byte byte() { long.byte; }
	Int int() { long.int; }
	Nat nat() { long.nat; }
	Word word() { long.word; }
	Long long() {
		if (type != Type:long)
			throw JsonAccessError("Expected integral number", this);
		numberVal;
	}
	Float float() { double.float; }
	Double double() {
		Double result;
		if (type == Type:long) {
			result = numberVal.double;
		} else if (type == Type:double) {
			Long copy = numberVal;
			asm { mov result, copy; }
		} else {
			throw JsonAccessError("Expected number", this);
		}
		result;
	}
	Str str() {
		unless (objectVal as Str)
			throw JsonAccessError("Expected string", this);
		objectVal;
	}
	JsonValue[] array() {
		unless (objectVal as JsonValue[])
			throw JsonAccessError("Expected array", this);
		objectVal;
	}
	Str->JsonValue object() {
		unless (objectVal as Str->JsonValue)
			throw JsonAccessError("Expected object", this);
		objectVal;
	}

	// Array/object interface.
	Nat count() {
		if (objectVal as JsonValue[])
			objectVal.count;
		else if (objectVal as Str->JsonValue)
			objectVal.count;
		else
			throw JsonAccessError("Expected array or object", this);
	}

	// Array interface.
	JsonValue [](Nat id) {
		array[id];
	}
	assign [](Nat id, JsonValue v) {
		// TODO: Maybe insert values and convert to array?
		array[id] = v;
	}
	JsonValue <<(JsonValue v) { push(v); this; }
	void push(JsonValue v) {
		if (type == Type:empty) {
			objectVal = JsonValue[];
			type = Type:array;
		}
		array.push(v);
	}
	void remove(Nat id) {
		array.remove(id);
	}

	// Object interface.
	JsonValue [](Str key) {
		object.get(key);
	}
	assign [](Str key, JsonValue v) { put(key, v); }
	JsonValue? at(Str key) {
		object.at(key);
	}
	void put(Str key, JsonValue v) {
		if (type == Type:empty) {
			objectVal = Str->JsonValue;
			type = Type:object;
		}
		object.put(key, v);
	}
	void remove(Str key) {
		object.remove(key);
	}

	// Todo: Custom iterator that gives keys/values in relevant format?

	// To string, providing format information. Pass `indent = 0` to skip linebreaks altogether.
	Str toS(Nat indent) {
		toS(indent, false);
	}

	// To string, provide the option to sort keys in objects.
	Str toS(Nat indent, Bool sortKeys) {
		StrBuf to;
		if (indent > 0)
			to.indentBy(" " * indent);
		this.toS(to, indent == 0, sortKeys);
		to.toS;
	}

	// To string, use the indentation already in strbuf.
	void toS(StrBuf to) {
		toS(to, false, false);
	}

	// To string, produce compact representation if asked to.
	void toS(StrBuf to, Bool compact, Bool sortKeys) {
		if (type == Type:bool) {
			if (numberVal != 0)
				to << "true";
			else
				to << "false";
		} else if (type == Type:long) {
			to << numberVal;
		} else if (type == Type:double) {
			Long copy = numberVal;
			Double out;
			asm { mov out, copy; }
			// TODO: Proper formatting of NAN and denormalized values.
			to << out;
		} else if (objectVal as Str) {
			formatStr(to, objectVal);
		} else if (objectVal as JsonValue[]) {
			to << "[";
			for (i, v in objectVal) {
				if (i > 0) {
					if (compact)
						to << ",";
					else
						to << ", ";
				}
				v.toS(to, compact, sortKeys);
			}
			to << "]";
		} else if (objectVal as Str->JsonValue) {
			to << "{";
			if (!compact) {
				to.indent();
			}

			if (sortKeys) {
				putOrdered(to, objectVal, compact);
			} else {
				putUnordered(to, objectVal, compact);
			}

			if (!compact) {
				to.dedent();
				to << "\n";
			}
			to << "}";
		} else {
			to << "null";
		}
	}

	// Helper to output unsorted keys.
	private void putOrdered(StrBuf to, Str->JsonValue obj, Bool compact) : static {
		Str[] order;
		order.reserve(obj.count);
		for (k, v in obj)
			order << k;

		order.sort();

		Bool first = true;
		for (k in order) {
			if (first) {
				if (!compact)
					to << "\n";
			} else {
				if (compact)
					to << ",";
				else
					to << ",\n";
			}
			first = false;

			formatStr(to, k);
			if (compact)
				to << ":";
			else
				to << ": ";
			obj.get(k).toS(to, compact, true);
		}
	}

	// Helper to output unsorted keys.
	private void putUnordered(StrBuf to, Str->JsonValue obj, Bool compact) : static {
		Bool first = true;
		for (k, v in obj) {
			if (first) {
				if (!compact)
					to << "\n";
			} else {
				if (compact)
					to << ",";
				else
					to << ",\n";
			}
			first = false;

			formatStr(to, k);
			if (compact)
				to << ":";
			else
				to << ": ";
			v.toS(to, compact, false);
		}
	}

	// The actual data:

	// Data type:
	private enum Type {
		empty,
		// Numeric types are stored in the 'numberVal' variable. Even floating point types.
		bool,
		long,
		double,
		// Stored in 'objectVal' variable.
		string,
		array,
		object,
	}
	private Type type;
	private Long numberVal;
	private Object? objectVal;
}


// Separate function to avoid recursing in the JsonValue.== function when comparing strings and integers.
private Bool compare(Long a, Long b) : inline {
	a == b;
}

private Bool compare(Bool a, Bool b) : inline {
	a == b;
}

private Bool compare(Str a, Str b) : inline {
	a == b;
}

// Helper for comparison.
private Bool compare(JsonValue[] a, JsonValue[] b) {
	if (a.count != b.count)
		return false;
	for (i, x in a)
		if (!(x == b[i]))
			return false;
	return true;
}

private Bool compare(Str->JsonValue a, Str->JsonValue b) {
	if (a.count != b.count)
		return false;
	for (k, x in a) {
		unless (y = b.at(k))
			return false;
		if (!(x == y))
			return false;
	}
	return true;
}

// Helper to format strings properly.
private void formatStr(StrBuf to, Str str) {
	to << '"';
	for (ch in str) {
		if (ch == '\\') {
			to << "\\\\";
		} else if (ch == '"') {
			to << "\\\"";
		} else if (ch == '\n') {
			to << "\\n";
		} else if (ch == '\r') {
			to << "\\r";
		} else if (ch.codepoint <= 0x1F) {
			// Need to be escaped, control character.
			to << "\\u" << hex(ch.codepoint, 4);
		} else if (ch.codepoint > 0x7F) {
			// Outside of ascii, so escape for good measure. Note that JS works in UTF-16.
			Nat leading = ch.utf16Leading();
			Nat trailing = ch.utf16Trailing();
			if (leading != 0)
				to << "\\u" << hex(leading, 4);
			to << "\\u" << hex(trailing, 4);
		} else {
			to << ch;
		}
	}
	to << '"';
}