package ru.dbotthepony.mc.otm.registry

import it.unimi.dsi.fastutil.objects.Object2ObjectLinkedOpenHashMap
import net.minecraft.Util
import net.minecraft.core.Registry
import net.minecraft.resources.ResourceKey
import net.minecraft.resources.ResourceLocation
import net.minecraft.world.item.DyeColor
import net.minecraftforge.eventbus.api.IEventBus
import net.minecraftforge.registries.IForgeRegistry
import net.minecraftforge.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.collect.SupplierMap
import java.util.LinkedList
import java.util.concurrent.FutureTask
import java.util.concurrent.locks.LockSupport
import java.util.function.Supplier
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: IForgeRegistry<R>, modId: String = OverdriveThatMatters.MOD_ID, allowParallel: Boolean = false) : this(registry.registryKey, modId, allowParallel)

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

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

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

		fun run(): List<Entry<R>> {
			if (serial) {
				LOGGER.debug("Serial registration of {} entries from {} for {}...", entries.size, modId, registry.location())
				val t = SystemTime()
				entries.forEach { it.initialize() }
				LOGGER.debug("Serial registration of {} entries from {} for {} took {} ms", entries.size, modId, registry.location(), t.millis)
			} else {
				LOGGER.debug("Parallel registration of {} entries from {} for {}...", entries.size, modId, registry.location())
				val t = SystemTime()
				val futures = LinkedList<FutureTask<R>>()
				entries.forEach { futures.add(FutureTask(it::initialize)) }
				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.location(), t.millis)
			}

			return entries
		}
	}

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

		init {
			check(allowRegistration) { "Unable to register new entries after RegisterEvent has been fired" }
			require(entries.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)
			}
		}

		val key = ResourceLocation(modId, name)
		private var _value: T? = null

		override val value: T get() {
			return _value ?: throw IllegalStateException("Trying to access $name 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(): T {
			check(_value == null) { "Already initialized $name of $registry!" }

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

			_value = getValue
			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)

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

			for (stage in stages) {
				stage.run().forEach {
					event.register(registry, it.key, it)
				}
			}
		}
	}

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

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