package ru.dbotthepony.mc.otm.client.render

import com.google.common.collect.ImmutableList
import com.mojang.blaze3d.vertex.BufferBuilder
import com.mojang.blaze3d.vertex.VertexConsumer
import com.mojang.blaze3d.vertex.VertexSorting
import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
import it.unimi.dsi.fastutil.objects.ReferenceArraySet
import net.minecraft.client.renderer.MultiBufferSource
import net.minecraft.client.renderer.RenderType
import net.minecraft.client.renderer.Sheets

private fun equals(existing: ImmutableList<RenderType>?, types: Array<out RenderType>): Boolean {
	if (types.isEmpty()) {
		return existing == null
	}

	if (existing == null) {
		return false
	}

	if (existing.size != types.size) {
		return false
	}

	for (i in types.indices) {
		if (types[i] !== existing[i]) {
			return false
		}
	}

	return true
}

private fun equals(existing: ImmutableList<RenderType>?, types: ImmutableList<RenderType>): Boolean {
	if (existing === types) {
		return true
	}

	if (types.isEmpty()) {
		return existing == null
	}

	if (existing == null) {
		return false
	}

	if (existing.size != types.size) {
		return false
	}

	for (i in types.indices) {
		if (types[i] !== existing[i]) {
			return false
		}
	}

	return true
}

/**
 * Consider this class internal, because it is almost as rigid at Mojang's [MultiBufferSource.BufferSource]
 *
 * Allows to batch OTM's geometry
 */
class DynamicBufferSource(val minimalInitialBufferSize: Int = 0, val maximalInitialBufferSize: Int = Int.MAX_VALUE, val vertexSorting: VertexSorting = VertexSorting.DISTANCE_TO_ORIGIN) : MultiBufferSource {
	init {
		require(minimalInitialBufferSize >= 0) { "Invalid minimal buffer size $minimalInitialBufferSize" }
		require(maximalInitialBufferSize >= minimalInitialBufferSize) { "Maximal buffer size $maximalInitialBufferSize must be greater or equal to minimal buffer size $minimalInitialBufferSize" }
	}

	private inner class State(
		val type: RenderType,
		var after: ImmutableList<RenderType>? = null,
		immutableAfter: Boolean = false,
		chained: Boolean = true
	) {
		val immutableAfter: Boolean
		val chained: Boolean

		init {
			if (
				type.name == "forge_text" ||
				type.name == "text_intensity" ||
				type.name == "forge_text_see_through" ||
				type.name.contains("font_") ||
				type.name.contains("text_") ||
				type.name.contains("_text") ||
				type.name.contains("_font")
			) {
				after = ImmutableList.of(Sheets.signSheet())
				this.immutableAfter = true
				this.chained = true
			} else {
				this.immutableAfter = immutableAfter
				this.chained = chained
			}
		}

		val dependants = ArrayList<State>(0)
		var priority = -1
		var dirty: Boolean = false
		val builder by lazy(LazyThreadSafetyMode.NONE) {
			BufferBuilder(type.bufferSize().coerceAtLeast(minimalInitialBufferSize).coerceAtMost(maximalInitialBufferSize))
		}
	}

	@Suppress("BooleanLiteralArgument")
	private val bufferList = ArrayList<State>().also {
		var next = State(Sheets.solidBlockSheet(), null, true, false)
		it.add(next)
		next = State(Sheets.cutoutBlockSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.bannerSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.translucentCullBlockSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.shieldSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.bedSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.shulkerBoxSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.signSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(Sheets.chestSheet(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.translucentNoCrumbling(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.armorGlint(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.armorEntityGlint(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.glint(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.glintDirect(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.glintTranslucent(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.entityGlint(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.entityGlintDirect(), ImmutableList.of(next.type), true, false)
		it.add(next)
		next = State(RenderType.waterMask(), ImmutableList.of(next.type), true, false)
		it.add(next)
	}

	private fun determineHeight(input: State, seen: MutableSet<State>) {
		if (!seen.add(input)) {
			throw IllegalStateException("Recursive dependencies detected when going through ${input.type}: ${seen.joinToString(", ") { it.type.toString() + "[${it.after?.joinToString(", ")}]" }}")
		}

		if (input.after == null) {
			input.priority = 0
		} else {
			for (type in input.after!!) {
				for (input2 in bufferList) {
					if (input2.type == type) {
						if (input2.priority == -1) {
							determineHeight(input2, seen)
						}

						input.priority = input.priority.coerceAtLeast(input2.priority + 1)

						if (input.chained) {
							input2.dependants.add(input)
						}

						break
					}
				}
			}

			if (input.priority == -1) {
				input.priority = 0
			}
		}
	}

	private fun sortBufferList() {
		for (state in bufferList) {
			state.priority = -1
			state.dependants.clear()
		}

		for (state in bufferList) {
			val seen = ReferenceArraySet<State>()
			determineHeight(state, seen)
		}

		bufferList.sortWith { a, b ->
			// most referenced buffers go first
			b.priority.compareTo(a.priority)
		}
	}

	private val buffers = Reference2ObjectOpenHashMap<RenderType, State>().also {
		for (v in bufferList) it[v.type] = v
	}

	init {
		sortBufferList()
	}

	override fun getBuffer(renderType: RenderType): VertexConsumer {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType)
			buffers[renderType] = getState
			bufferList.add(getState)
			getState.dirty = true
			return getState.builder.also {
				it.begin(renderType.mode(), renderType.format())
			}
		}

		if (!getState.dirty) {
			getState.dirty = true
			getState.builder.begin(renderType.mode(), renderType.format())
		}

		return getState.builder
	}

	fun getBuffer(renderType: RenderType, after: RenderType): VertexConsumer {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType, ImmutableList.of(after))
			buffers[renderType] = getState
			bufferList.add(getState)
			getState.dirty = true
			sortBufferList()
			return getState.builder.also {
				it.begin(renderType.mode(), renderType.format())
			}
		}

		if (!getState.immutableAfter && getState.after?.getOrNull(0) !== after) {
			getState.after = ImmutableList.of(after)
			sortBufferList()
		}

		if (!getState.dirty) {
			getState.dirty = true
			getState.builder.begin(renderType.mode(), renderType.format())
		}

		return getState.builder
	}

	fun getBuffer(renderType: RenderType, vararg after: RenderType): VertexConsumer {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType, ImmutableList.copyOf(after))
			buffers[renderType] = getState
			bufferList.add(getState)
			getState.dirty = true
			sortBufferList()
			return getState.builder.also {
				it.begin(renderType.mode(), renderType.format())
			}
		}

		if (!getState.immutableAfter && !equals(getState.after, after)) {
			getState.after = ImmutableList.copyOf(after)
			sortBufferList()
		}

		if (!getState.dirty) {
			getState.dirty = true
			getState.builder.begin(renderType.mode(), renderType.format())
		}

		return getState.builder
	}

	fun getBuffer(renderType: RenderType, after: ImmutableList<RenderType>): VertexConsumer {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType, after)
			buffers[renderType] = getState
			bufferList.add(getState)
			getState.dirty = true
			sortBufferList()
			return getState.builder.also {
				it.begin(renderType.mode(), renderType.format())
			}
		}

		if (!getState.immutableAfter && !equals(getState.after, after)) {
			getState.after = after
			sortBufferList()
		}

		if (!getState.dirty) {
			getState.dirty = true
			getState.builder.begin(renderType.mode(), renderType.format())
		}

		return getState.builder
	}

	fun registerType(renderType: RenderType, after: RenderType) {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType, ImmutableList.of(after), immutableAfter = true)
			buffers[renderType] = getState
			bufferList.add(getState)
			sortBufferList()
		} else if (!getState.immutableAfter && getState.after?.getOrNull(0) !== after) {
			getState.after = ImmutableList.of(after)
			sortBufferList()
		}
	}

	fun registerType(renderType: RenderType, vararg after: RenderType) {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType, ImmutableList.copyOf(after), immutableAfter = true)
			buffers[renderType] = getState
			bufferList.add(getState)
			sortBufferList()
		} else if (!getState.immutableAfter && !equals(getState.after, after)) {
			getState.after = ImmutableList.copyOf(after)
			sortBufferList()
		}
	}

	fun registerType(renderType: RenderType, after: ImmutableList<RenderType>) {
		var getState = buffers[renderType]

		if (getState == null) {
			getState = State(renderType, after, immutableAfter = true)
			buffers[renderType] = getState
			bufferList.add(getState)
			sortBufferList()
		} else if (!getState.immutableAfter && !equals(getState.after, after)) {
			getState.after = after
			sortBufferList()
		}
	}

	fun endBatch() {
		for (state in bufferList) {
			if (state.dirty) {
				state.dirty = false
				state.type.end(state.builder, vertexSorting)
			}
		}
	}

	fun endBatch(type: RenderType) {
		val state = buffers[type] ?: return

		if (state.dirty) {
			state.dirty = false
			type.end(state.builder, vertexSorting)
		}
	}

	fun endBatchChain(type: RenderType) {
		val state = buffers[type] ?: return

		if (state.dirty) {
			state.dirty = false
			type.end(state.builder, vertexSorting)
		}

		for (ustate in state.dependants) {
			endBatchChain(ustate)
		}
	}

	private fun endBatchChain(state: State) {
		if (state.dirty) {
			state.dirty = false
			state.type.end(state.builder, vertexSorting)
		}

		for (ustate in state.dependants) {
			endBatchChain(ustate)
		}
	}

	companion object {
		@JvmField
		val GUI = DynamicBufferSource(maximalInitialBufferSize = 2 shl 8, vertexSorting = VertexSorting.ORTHOGRAPHIC_Z)
		@JvmField
		val WORLD = DynamicBufferSource(minimalInitialBufferSize = 2 shl 12)

		/**
		 * Called from LevelRenderer
		 */
		@JvmStatic
		fun renderHook(type: RenderType) {
			WORLD.endBatchChain(type)
		}

		/**
		 * Called from LevelRenderer
		 */
		@JvmStatic
		fun renderHook() {
			WORLD.endBatch()
		}
	}
}
