package moe.nea.firmament.util.mc

import com.google.gson.Gson
import com.google.gson.JsonArray
import com.google.gson.JsonElement
import com.google.gson.JsonPrimitive
import com.mojang.brigadier.StringReader
import com.mojang.brigadier.arguments.ArgumentType
import com.mojang.brigadier.arguments.StringArgumentType
import com.mojang.serialization.JsonOps
import kotlin.jvm.optionals.getOrNull
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.Tag
import net.minecraft.nbt.ListTag
import net.minecraft.nbt.NbtOps
import net.minecraft.nbt.StringTag
import moe.nea.firmament.util.Base64Util

class NbtPrism(val path: List<String>) {
	companion object {
		fun fromElement(path: JsonElement): NbtPrism? {
			if (path is JsonArray) {
				return NbtPrism(path.map { (it as JsonPrimitive).asString })
			} else if (path is JsonPrimitive && path.isString) {
				return NbtPrism(path.asString.split("."))
			}
			return null
		}
	}

	object Argument : ArgumentType<NbtPrism> {
		override fun parse(reader: StringReader): NbtPrism? {
			return fromElement(JsonPrimitive(StringArgumentType.string().parse(reader)))
		}

		override fun getExamples(): Collection<String?>? {
			return listOf("some.nbt.path", "some.other.*", "some.path.*json.in.a.json.string")
		}
	}

	override fun toString(): String {
		return "Prism($path)"
	}

	fun access(root: Tag): Collection<Tag> {
		var rootSet = mutableListOf(root)
		var switch = mutableListOf<Tag>()
		for (pathSegment in path) {
			if (pathSegment == ".") continue
			if (pathSegment != "*" && pathSegment.startsWith("*")) {
				if (pathSegment == "*json") {
					for (element in rootSet) {
						val eString = element.asString().getOrNull() ?: continue
						val element = Gson().fromJson(eString, JsonElement::class.java)
						switch.add(JsonOps.INSTANCE.convertTo(NbtOps.INSTANCE, element))
					}
				} else if (pathSegment == "*base64") {
					for (element in rootSet) {
						val string = element.asString().getOrNull() ?: continue
						switch.add(StringTag.valueOf(Base64Util.decodeString(string)))
					}
				}
			}
			for (element in rootSet) {
				if (element is ListTag) {
					if (pathSegment == "*")
						switch.addAll(element)
					val index = pathSegment.toIntOrNull() ?: continue
					if (index !in element.indices) continue
					switch.add(element[index])
				}
				if (element is CompoundTag) {
					if (pathSegment == "*")
						element.keySet().mapTo(switch) { element.get(it)!! }
					switch.add(element.get(pathSegment) ?: continue)
				}
			}
			val temp = switch
			switch = rootSet
			rootSet = temp
			switch.clear()
		}
		return rootSet
	}
}