package ru.dbotthepony.mc.otm.core.util

import com.google.common.collect.ImmutableList
import com.google.gson.JsonArray
import com.google.gson.JsonElement
import com.google.gson.JsonNull
import com.google.gson.JsonObject
import com.google.gson.JsonParseException
import com.google.gson.JsonPrimitive
import com.google.gson.JsonSyntaxException
import net.minecraft.nbt.NbtAccounter
import ru.dbotthepony.mc.otm.core.collect.stream
import java.io.InputStream
import java.io.OutputStream
import java.util.function.Predicate

private enum class BinaryElementType(private val predicate: Predicate<JsonElement>) : Predicate<JsonElement> by predicate {
	NULL({ it is JsonNull }) {
		override fun write(stream: OutputStream, element: JsonElement) {
		}

		override fun read(stream: InputStream): JsonElement {
			return JsonNull.INSTANCE
		}
	},

	DOUBLE({ it is JsonPrimitive && it.isNumber && (it.asNumber is Double || it.asNumber is Float) }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			stream.writeDouble(element.asDouble)
		}

		override fun read(stream: InputStream): JsonElement {
			return JsonPrimitive(stream.readDouble())
		}
	},

	BOOLEAN({ it is JsonPrimitive && it.isBoolean }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			stream.write(if (element.asBoolean) 1 else 0)
		}

		override fun read(stream: InputStream): JsonElement {
			return JsonPrimitive(stream.read() > 0)
		}
	},

	INT({ it is JsonPrimitive && it.isNumber }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			val it = element.asNumber

			if (it is Int || it is Long || it is Short || it is Byte)
				stream.writeVarLongLE(element.asLong)
			else
				throw IllegalArgumentException("Unknown number type: ${it::class.qualifiedName}")
		}

		override fun read(stream: InputStream): JsonElement {
			return JsonPrimitive(stream.readVarLongLE())
		}
	},

	STRING({ it is JsonPrimitive && it.isString }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			stream.writeBinaryString(element.asString)
		}

		override fun read(stream: InputStream): JsonElement {
			return JsonPrimitive(stream.readBinaryString())
		}
	},

	ARRAY({ it is JsonArray }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			writeArray(stream, element, OutputStream::writeBinaryJson)
		}

		override fun read(stream: InputStream): JsonElement {
			return readArray(stream, InputStream::readBinaryJson)
		}
	},

	OBJECT({ it is JsonObject }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			element as JsonObject
			stream.writeVarIntLE(element.size())

			for ((k, v) in element.entrySet()) {
				stream.writeBinaryString(k)
				stream.writeBinaryJson(v)
			}
		}

		override fun read(stream: InputStream): JsonElement {
			val count = stream.readVarIntLE()

			if (count == 0) return JsonObject()
			if (count < 0) throw JsonSyntaxException("Tried to read json object with $count elements in it")

			val build = JsonObject()

			for (i in 0 until count) {
				val key = try {
					stream.readBinaryString()
				} catch(err: Throwable) {
					throw JsonSyntaxException("Reading json object at $i", err)
				}

				try {
					build.add(key, stream.readBinaryJson())
				} catch(err: Throwable) {
					throw JsonSyntaxException("Reading json object at $i with name $key", err)
				}
			}

			return build
		}
	},

	DOUBLE_ARRAY({ it is JsonArray && it.stream().allMatch { DOUBLE.test(it) } }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			writeArray(stream, element, DOUBLE::write)
		}

		override fun read(stream: InputStream): JsonElement {
			return readArray(stream, DOUBLE::read)
		}
	},

	INT_ARRAY({ it is JsonArray && it.stream().allMatch { INT.test(it) } }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			writeArray(stream, element, INT::write)
		}

		override fun read(stream: InputStream): JsonElement {
			return readArray(stream, INT::read)
		}
	},

	BOOLEAN_ARRAY({ it is JsonArray && it.stream().allMatch { BOOLEAN.test(it) } }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			writeArray(stream, element, BOOLEAN::write)
		}

		override fun read(stream: InputStream): JsonElement {
			return readArray(stream, BOOLEAN::read)
		}
	},

	NULL_ARRAY({ it is JsonArray && it.stream().allMatch { NULL.test(it) } }) {
		override fun write(stream: OutputStream, element: JsonElement) {
			writeArray(stream, element, NULL::write)
		}

		override fun read(stream: InputStream): JsonElement {
			return readArray(stream, NULL::read)
		}
	};

	protected fun writeArray(stream: OutputStream, element: JsonElement, writer: (OutputStream, JsonElement) -> Unit) {
		element as JsonArray
		stream.writeVarIntLE(element.size())

		for (v in element) {
			writer.invoke(stream, v)
		}
	}

	protected fun readArray(stream: InputStream, reader: (InputStream) -> JsonElement): JsonArray {
		val count = stream.readVarIntLE()

		if (count == 0) return JsonArray()
		if (count < 0) throw JsonSyntaxException("Tried to read json array with $count elements in it")

		return JsonArray(count).also {
			for (i in 0 until count) it.add(reader.invoke(stream))
		}
	}

	abstract fun write(stream: OutputStream, element: JsonElement)
	abstract fun read(stream: InputStream): JsonElement

	companion object {
		val cached: ImmutableList<BinaryElementType> = ImmutableList.of(
			NULL, DOUBLE, BOOLEAN, INT, STRING, OBJECT,
			DOUBLE_ARRAY, INT_ARRAY, BOOLEAN_ARRAY, ARRAY)

	}
}

/**
 * Writes json in efficient binary form to stream
 */
fun OutputStream.writeBinaryJson(element: JsonElement) {
	for (writer in BinaryElementType.cached) {
		if (writer.test(element)) {
			write(writer.ordinal + 1)
			writer.write(this, element)
			return
		}
	}

	throw IllegalArgumentException("Unknown element type: ${element::class.qualifiedName}")
}

/**
 * Reads json in binary form from stream
 */
fun InputStream.readBinaryJson(): JsonElement {
	val id = read() - 1
	val reader = BinaryElementType.entries.getOrNull(id) ?: throw JsonParseException("Unknown element type ${id + 1}")
	return reader.read(this)
}
