package ru.dbotthepony.mc.otm.registry

import com.mojang.datafixers.util.Either
import it.unimi.dsi.fastutil.objects.Object2ObjectLinkedOpenHashMap
import net.minecraft.Util
import net.minecraft.core.Holder
import net.minecraft.core.HolderOwner
import net.minecraft.core.Registry
import net.minecraft.resources.ResourceKey
import net.minecraft.resources.ResourceLocation
import net.minecraft.tags.TagKey
import net.minecraft.world.item.DyeColor
import net.neoforged.api.distmarker.Dist
import net.neoforged.bus.api.IEventBus
import net.neoforged.fml.loading.FMLEnvironment
import net.neoforged.neoforge.capabilities.RegisterCapabilitiesEvent
import net.neoforged.neoforge.client.extensions.common.RegisterClientExtensionsEvent
import net.neoforged.neoforge.registries.RegisterEvent
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.mc.otm.OverdriveThatMatters
import ru.dbotthepony.mc.otm.SystemTime
import ru.dbotthepony.mc.otm.core.ResourceLocation
import ru.dbotthepony.mc.otm.core.collect.SupplierMap
import java.util.*
import java.util.concurrent.FutureTask
import java.util.concurrent.locks.LockSupport
import java.util.function.Predicate
import java.util.function.Supplier
import java.util.stream.Stream
import kotlin.reflect.KProperty

/**
 * DeferredRegister which allows parallel initialization
 */
class MDeferredRegister<R : Any>(val registry: ResourceKey<out Registry<R>>, val modId: String = OverdriveThatMatters.MOD_ID, val allowParallel: Boolean = false) {
	constructor(registry: Registry<R>, modId: String = OverdriveThatMatters.MOD_ID, allowParallel: Boolean = false) : this(registry.key(), modId, allowParallel)

	private val entriesInternal = Object2ObjectLinkedOpenHashMap<ResourceLocation, Entry<R>>()
	private var allowRegistration = true
	private var eventRegistered = false
	private val stages = ArrayList<RegistrationStage>()

	val entries: Map<ResourceLocation, Entry<R>> = Collections.unmodifiableMap(entriesInternal)

	private inner class RegistrationStage(val serial: Boolean) {
		private val entries = ArrayList<Entry<R>>()

		fun add(entry: Entry<R>) {
			entries.add(entry)
		}

		fun run(registry: Registry<R>) {
			check(registry.key() == this@MDeferredRegister.registry) { "Hey." }

			if (serial) {
				LOGGER.debug("Serial registration of {} entries from {} for {}...", entries.size, modId, registry.key().location())
				val t = SystemTime()
				entries.forEach { it.initialize(registry) }
				LOGGER.debug("Serial registration of {} entries from {} for {} took {} ms", entries.size, modId, registry.key().location(), t.millis)
			} else {
				LOGGER.debug("Parallel registration of {} entries from {} for {}...", entries.size, modId, registry.key().location())
				val t = SystemTime()
				val futures = LinkedList<FutureTask<R>>()
				entries.forEach { futures.add(FutureTask { it.initialize(registry) }) }
				futures.forEach { Util.backgroundExecutor().submit(it) }

				while (futures.isNotEmpty()) {
					futures.removeIf { it.run(); it.isDone } // help executor threads to minimize thread stalling
					// especially true when executor threads are busy with work created by other mods
					// so we keep worst-case scenario where every registry entry is initialized and registered serially as it was before
					LockSupport.parkNanos(1_000_000L)
				}

				// memory barrier to avoid specifying _value as @Volatile
				LockSupport.parkNanos(1_000_000L)

				LOGGER.debug("Parallel registration of {} entries from {} for {} took {} ms", entries.size, modId, registry.key().location(), t.millis)
			}
		}
	}

	inner class Entry<out T : R>(name: ResourceLocation, private val factory: Supplier<T>, serial: Boolean) : Supplier<@UnsafeVariance T>, Lazy<T>, Holder<R> {
		constructor(name: String, factory: Supplier<T>, serial: Boolean) : this(ResourceLocation(modId, name), factory, serial)
		constructor(name: String, factory: Supplier<T>) : this(ResourceLocation(modId, name), factory, false)
		constructor(name: ResourceLocation, factory: Supplier<T>) : this(name, factory, false)

		init {
			check(allowRegistration) { "Unable to register new entries after RegisterEvent has been fired" }
			require(entriesInternal.put(name, this) == null) { "Duplicate entry $name for registry $registry" }

			val actualSerial = if (allowParallel) serial else true

			if (stages.isEmpty() || stages.last().serial != actualSerial) {
				val stage = RegistrationStage(actualSerial)
				stage.add(this)
				stages.add(stage)
			} else {
				stages.last().add(this)
			}
		}

		@JvmField
		val key: ResourceKey<R> = ResourceKey.create(registry, name)
		private var parentHolder: Holder<R>? = null

		override fun value(): T {
			return value
		}

		override fun isBound(): Boolean {
			return _value != null
		}

		override fun `is`(p_205713_: ResourceLocation): Boolean {
			return p_205713_ == key.location()
		}

		override fun `is`(p_205712_: ResourceKey<R>): Boolean {
			return p_205712_ == key
		}

		override fun `is`(p_205711_: Predicate<ResourceKey<R>>): Boolean {
			return p_205711_.test(key)
		}

		override fun `is`(p_205705_: TagKey<R>): Boolean {
			return parentHolder?.`is`(p_205705_) ?: false
		}

		@Deprecated("Deprecated in Java")
		override fun `is`(p_316447_: Holder<R>): Boolean {
			return p_316447_.`is`(key)
		}

		override fun tags(): Stream<TagKey<R>> {
			return parentHolder?.tags() ?: Stream.empty()
		}

		override fun unwrap(): Either<ResourceKey<R>, R> {
			return Either.left(key)
		}

		override fun unwrapKey(): Optional<ResourceKey<R>> {
			return Optional.of(key)
		}

		override fun kind(): Holder.Kind {
			return Holder.Kind.REFERENCE
		}

		override fun canSerializeIn(p_255833_: HolderOwner<R>): Boolean {
			return parentHolder?.canSerializeIn(p_255833_) ?: false
		}

		private var _value: T? = null

		override val value: T get() {
			return _value ?: throw IllegalStateException("Trying to access ${key.location()} of $registry before it is initialized")
		}

		override fun isInitialized(): Boolean {
			return _value != null
		}

		override fun get(): T {
			return value
		}

		operator fun getValue(receiver: Any?, property: KProperty<*>): T {
			return value
		}

		fun initialize(registry: Registry<R>): T {
			check(_value == null) { "Already initialized ${key.location()} of $registry!" }

			val compute = try {
				factory.get()
			} catch (err: Throwable) {
				throw RuntimeException("Unable to initialize registry entry ${key.location()} of $registry", err)
			}

			_value = compute
			parentHolder = Registry.registerForHolder(registry, key, compute)
			return value
		}
	}

	fun <T : R> colored(prefix: String, factory: (color: DyeColor) -> T): Map<DyeColor, T> {
		return SupplierMap(MRegistry.DYE_ORDER.map { it to Entry(prefix + "_" + it.name.lowercase()) { factory.invoke(it) } })
	}

	fun <T : R> coloredWithBase(prefix: String, factory: (color: DyeColor?) -> T): Map<DyeColor?, T> {
		val values = ArrayList<Pair<DyeColor?, Supplier<T>>>()

		values.add(null to Entry(prefix) { factory.invoke(null) })
		MRegistry.DYE_ORDER.forEach { values.add(it to Entry(prefix + "_" + it.name.lowercase()) { factory.invoke(it) }) }

		return SupplierMap(values)
	}

	fun <T : R> register(name: String, factory: Supplier<T>) = Entry(name, factory)
	fun <T : R> register(name: String, factory: Supplier<T>, serial: Boolean) = Entry(name, factory, serial)
	fun <T : R> register(name: String, serial: Boolean, factory: Supplier<T>) = Entry(name, factory, serial)

	fun <T : R> register(name: ResourceLocation, factory: Supplier<T>) = Entry(name, factory)
	fun <T : R> register(name: ResourceLocation, factory: Supplier<T>, serial: Boolean) = Entry(name, factory, serial)
	fun <T : R> register(name: ResourceLocation, serial: Boolean, factory: Supplier<T>) = Entry(name, factory, serial)

	fun <T : R> register(name: ResourceLocation, value: T) = Entry(name) { value }

	private fun onRegisterEvent(event: RegisterEvent) {
		if (event.registryKey == registry) {
			allowRegistration = false

			for (stage in stages) {
				stage.run(event.registry as Registry<R>)
			}
		}
	}

	fun register(bus: IEventBus) {
		check(!eventRegistered) { "Already registered!" }
		eventRegistered = true
		bus.addListener(this::onRegisterEvent)
		bus.addListener(this::registerCapabilities)

		if (FMLEnvironment.dist == Dist.CLIENT) {
			bus.addListener(this::registerClientExtensions)
		}
	}

	private fun registerCapabilities(event: RegisterCapabilitiesEvent) {
		entriesInternal.values.forEach {
			val value = it.value

			if (value is CapabilitiesRegisterListener) {
				value.registerCapabilities(event)
			}
		}
	}

	private fun registerClientExtensions(event: RegisterClientExtensionsEvent) {
		entriesInternal.values.forEach {
			val value = it.value

			if (value is ClientExtensionsListener) {
				value.registerClientExtensions(event)
			}
		}
	}

	companion object {
		private val LOGGER = LogManager.getLogger()
	}
}
