package ru.dbotthepony.mc.otm.core.nbt

import com.google.gson.JsonElement
import it.unimi.dsi.fastutil.io.FastByteArrayInputStream
import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
import net.minecraft.nbt.ByteArrayTag
import net.minecraft.nbt.ByteTag
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.DoubleTag
import net.minecraft.nbt.FloatTag
import net.minecraft.nbt.IntArrayTag
import net.minecraft.nbt.IntTag
import net.minecraft.nbt.ListTag
import net.minecraft.nbt.LongArrayTag
import net.minecraft.nbt.LongTag
import net.minecraft.nbt.NbtAccounter
import net.minecraft.nbt.NbtUtils
import net.minecraft.nbt.NumericTag
import net.minecraft.nbt.ShortTag
import net.minecraft.nbt.StringTag
import net.minecraft.nbt.Tag
import net.minecraft.world.item.ItemStack
import ru.dbotthepony.mc.otm.core.util.readJson
import ru.dbotthepony.mc.otm.core.util.writeJson
import java.util.UUID

operator fun CompoundTag.set(index: String, value: Tag) = 				put(index, value)
operator fun CompoundTag.set(index: String, value: Int) = 				putInt(index, value)
operator fun CompoundTag.set(index: String, value: Byte) = 				putByte(index, value)
operator fun CompoundTag.set(index: String, value: Short) = 			putShort(index, value)
operator fun CompoundTag.set(index: String, value: Long) = 				putLong(index, value)
operator fun CompoundTag.set(index: String, value: Float) = 			putFloat(index, value)
operator fun CompoundTag.set(index: String, value: Double) = 			putDouble(index, value)
operator fun CompoundTag.set(index: String, value: String) = 			putString(index, value)
operator fun CompoundTag.set(index: String, value: Boolean) = 			putBoolean(index, value)
operator fun CompoundTag.set(index: String, value: ByteArray) = 		putByteArray(index, value)
operator fun CompoundTag.set(index: String, value: IntArray) = 			putIntArray(index, value)
operator fun CompoundTag.set(index: String, value: LongArray) = 		putLongArray(index, value)

operator fun CompoundTag.set(index: String, value: UUID) = 				putUUID(index, value)

fun CompoundTag.contains(vararg keys: String): Boolean {
	for (key in keys) {
		if (!contains(key)) {
			return false
		}
	}

	return true
}

inline fun <R, reified T : Tag?> CompoundTag.map(key: String, consumer: (T) -> R): R? {
	val tag = get(key)

	if (tag is T) {
		return consumer(tag)
	}

	return null
}

inline fun <R, reified T : Tag> CompoundTag.mapPresent(key: String, consumer: (T) -> R): R? {
	val tag = get(key)

	if (tag is T) {
		return consumer(tag)
	}

	return null
}

fun <T> CompoundTag.mapString(index: String, mapper: (String) -> T, orElse: T): T {
	val tag = this[index] as? StringTag ?: return orElse
	return mapper.invoke(tag.asString)
}

fun CompoundTag.getItemStack(key: String): ItemStack = map(key, ItemStack::of) ?: ItemStack.EMPTY

@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getByteList(key: String): MutableList<ByteTag> = getList(key, Tag.TAG_BYTE.toInt()) as MutableList<ByteTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getShortList(key: String): MutableList<ShortTag> = getList(key, Tag.TAG_SHORT.toInt()) as MutableList<ShortTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getIntList(key: String): MutableList<IntTag> = getList(key, Tag.TAG_INT.toInt()) as MutableList<IntTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getLongList(key: String): MutableList<LongTag> = getList(key, Tag.TAG_LONG.toInt()) as MutableList<LongTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getFloatList(key: String): MutableList<FloatTag> = getList(key, Tag.TAG_FLOAT.toInt()) as MutableList<FloatTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getDoubleList(key: String): MutableList<DoubleTag> = getList(key, Tag.TAG_DOUBLE.toInt()) as MutableList<DoubleTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getByteArrayList(key: String): MutableList<ByteArrayTag> = getList(key, Tag.TAG_BYTE_ARRAY.toInt()) as MutableList<ByteArrayTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getStringList(key: String): MutableList<StringTag> = getList(key, Tag.TAG_STRING.toInt()) as MutableList<StringTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getListList(key: String): MutableList<ListTag> = getList(key, Tag.TAG_LIST.toInt()) as MutableList<ListTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getCompoundList(key: String): MutableList<CompoundTag> = getList(key, Tag.TAG_COMPOUND.toInt()) as MutableList<CompoundTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getIntArrayList(key: String): MutableList<IntArrayTag> = getList(key, Tag.TAG_INT_ARRAY.toInt()) as MutableList<IntArrayTag>
@Suppress("unchecked_cast") // type is checked inside getList
fun CompoundTag.getLongArrayList(key: String): MutableList<LongArrayTag> = getList(key, Tag.TAG_LONG_ARRAY.toInt()) as MutableList<LongArrayTag>

fun CompoundTag.getUUIDSafe(key: String): UUID? {
	val value = this[key] as? IntArrayTag ?: return null
	return NbtUtils.loadUUID(value)
}

fun CompoundTag.putJson(key: String, json: JsonElement) {
	val bytes = FastByteArrayOutputStream()
	bytes.writeJson(json)
	putByteArray(key, bytes.array.copyOfRange(0, bytes.length))
}

fun CompoundTag.getJson(key: String, sizeLimit: NbtAccounter = NbtAccounter.UNLIMITED): JsonElement? {
	val bytes = getByteArray(key)

	if (bytes.isEmpty())
		return null

	return FastByteArrayInputStream(bytes).readJson(sizeLimit)
}

fun CompoundTag.getBoolean(index: String, orElse: Boolean): Boolean {
	return (this[index] as? NumericTag)?.asInt?.let { it > 0 } ?: orElse
}
