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

import com.google.common.collect.ImmutableList
import com.mojang.blaze3d.vertex.BufferBuilder
import com.mojang.blaze3d.vertex.ByteBufferBuilder
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
import net.minecraft.client.resources.model.ModelBakery
import net.neoforged.fml.ModLoader
import net.neoforged.neoforge.client.event.RegisterRenderBuffersEvent
import ru.dbotthepony.mc.otm.client.render.blockentity.BlackHoleGeneratorRenderer
import java.io.Closeable

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, Closeable {
	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,
		private val byteBuffer: ByteBufferBuilder = ByteBufferBuilder(type.bufferSize().coerceIn(minimalInitialBufferSize, maximalInitialBufferSize))
	) {
		val immutableAfter: Boolean
		val chained: Boolean

		init {
			if (
				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
			}

			require(buffers.put(type, this) == null) { "Duplicate render type: $type" }
			bufferList.add(this)
		}

		val dependants = ArrayList<State>(0)
		var priority = -1

		private var buffer: BufferBuilder? = null

		fun begin(): BufferBuilder {
			var buffer = buffer

			if (buffer != null) {
				return buffer
			}

			buffer = BufferBuilder(byteBuffer, type.mode, type.format)
			this.buffer = buffer
			return buffer
		}

		fun finish() {
			val built = buffer?.build()

			if (built != null) {
				if (type.sortOnUpload()) {
					built.sortQuads(byteBuffer, vertexSorting)
				}

				type.draw(built)
			}

			this.buffer = null
		}

		fun finishChain() {
			finish()
			dependants.forEach { it.finishChain() }
		}

		fun close() {
			byteBuffer.close()
		}
	}

	private val buffers = Reference2ObjectOpenHashMap<RenderType, State>()
	private val bufferList = ArrayList<State>()

	init {
		var next = State(Sheets.solidBlockSheet(), null, immutableAfter = true, chained = false)
		next = State(Sheets.cutoutBlockSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(Sheets.bannerSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(Sheets.translucentCullBlockSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)

		next = State(Sheets.shieldSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(Sheets.bedSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(Sheets.shulkerBoxSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(Sheets.signSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		//next = State(Sheets.hangingSignSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)

		next = State(Sheets.chestSheet(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(CHART_RENDER_TYPE, ImmutableList.of(next.type), immutableAfter = true, chained = true)
		next = State(CHART_RENDER_TYPE2, ImmutableList.of(next.type), immutableAfter = true, chained = true)

		next = State(rectRenderType(true, false, true), ImmutableList.of(next.type), immutableAfter = true, chained = true)
		next = State(rectRenderType(true, false, false), ImmutableList.of(next.type), immutableAfter = true, chained = true)
		next = State(rectRenderType(false, false, true), ImmutableList.of(next.type), immutableAfter = true, chained = true)
		next = State(rectRenderType(false, false, false), ImmutableList.of(next.type), immutableAfter = true, chained = true)

		next = State(RenderType.armorEntityGlint(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(RenderType.glint(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(RenderType.glintTranslucent(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(RenderType.entityGlint(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(RenderType.entityGlintDirect(), ImmutableList.of(next.type), immutableAfter = true, chained = false)
		next = State(RenderType.waterMask(), ImmutableList.of(next.type), immutableAfter = true, chained = false)

		next = State(MULTIBLOCK_STATUS_RENDER_TYPE, ImmutableList.of(next.type), immutableAfter = true, chained = true)
		next = State(rectRenderType(true, true, true), ImmutableList.of(next.type), immutableAfter = true, chained = true)
		next = State(rectRenderType(true, true, false), ImmutableList.of(next.type), immutableAfter = true, chained = true)

		ModelBakery.DESTROY_TYPES.forEach {
			next = State(it, ImmutableList.of(next.type), immutableAfter = true, chained = false)
		}

		val modBuffers = LinkedHashMap<RenderType, ByteBufferBuilder>()
		ModLoader.postEvent(RegisterRenderBuffersEvent(modBuffers))

		for ((type, buffer) in modBuffers) {
			// TODO: determine whenever other mods render types should be chained
			next = State(type, ImmutableList.of(next.type), immutableAfter = false, chained = true, buffer)
		}

		sortBufferList()
	}

	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)
		}
	}

	init {
		sortBufferList()
	}

	override fun getBuffer(renderType: RenderType): VertexConsumer {
		return (buffers[renderType] ?: State(renderType)).begin()
	}

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

		if (getState == null) {
			getState = State(renderType, ImmutableList.of(after))
			sortBufferList()
			return getState.begin()
		}

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

		return getState.begin()
	}

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

		if (getState == null) {
			getState = State(renderType, ImmutableList.copyOf(after))
			sortBufferList()
			return getState.begin()
		}

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

		return getState.begin()
	}

	fun registerType(renderType: RenderType, after: RenderType) {
		return registerType(renderType, ImmutableList.of(after))
	}

	fun registerType(renderType: RenderType, vararg after: RenderType) {
		return registerType(renderType, ImmutableList.copyOf(after))
	}

	fun registerType(renderType: RenderType, after: List<RenderType>) {
		val afterList = ImmutableList.copyOf(after)
		var getState = buffers[renderType]

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

	fun endBatch() {
		bufferList.forEach { it.finish() }
	}

	fun endBatch(type: RenderType) {
		buffers[type]?.finish()
	}

	fun endBatchChain(type: RenderType) {
		buffers[type]?.finishChain()
	}

	override fun close() {
		bufferList.forEach { it.close() }
	}

	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()
		}
	}
}
