KStarbound/src/main/kotlin/ru/dbotthepony/kstarbound/Registry.kt

package ru.dbotthepony.kstarbound

import com.google.gson.JsonElement
import com.google.gson.JsonNull
import com.google.gson.JsonObject
import it.unimi.dsi.fastutil.ints.Int2ObjectFunction
import it.unimi.dsi.fastutil.ints.Int2ObjectMap
import it.unimi.dsi.fastutil.ints.Int2ObjectMaps
import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap
import it.unimi.dsi.fastutil.longs.LongOpenHashSet
import it.unimi.dsi.fastutil.objects.Object2ObjectFunction
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.kommons.util.Either
import ru.dbotthepony.kommons.util.KOptional
import ru.dbotthepony.kommons.util.XXHash64
import ru.dbotthepony.kstarbound.util.limit
import java.util.Collections
import java.util.concurrent.locks.ReentrantReadWriteLock
import java.util.function.Supplier
import kotlin.collections.set
import kotlin.concurrent.read
import kotlin.concurrent.write

class Registry<T : Any>(val name: String, val storeJson: Boolean = true) {
	private val keysInternal = HashMap<String, Impl>()
	private val idsInternal = Int2ObjectOpenHashMap<Impl>()
	private val keyRefs = HashMap<String, RefImpl>()
	private val idRefs = Int2ObjectOpenHashMap<RefImpl>()

	private val lock = ReentrantReadWriteLock()

	private var hasBeenValidated = false

	// idiot-proof miss lookup. Surely, it will cause some entries to be never logged
	// if they are missing, but at least if malicious actor spams with long-ass invalid data
	// it won't explode memory usage of server
	private val loggedMisses = LongOpenHashSet()

	val keys: Map<String, Entry<T>> = Collections.unmodifiableMap(keysInternal)
	val ids: Int2ObjectMap<out Entry<T>> = Int2ObjectMaps.unmodifiable(idsInternal)

	sealed class Ref<T : Any> : Supplier<Entry<T>?> {
		abstract val key: Either<String, Int>
		abstract val entry: Entry<T>?
		abstract val registry: Registry<T>

		val isPresent: Boolean
			get() = value != null

		val isEmpty: Boolean
			get() = value == null

		val value: T?
			get() = entry?.value

		final override fun get(): Entry<T>? {
			return entry
		}
	}

	sealed class Entry<T : Any> : Supplier<T> {
		abstract val key: String
		abstract val id: Int?
		abstract val value: T
		abstract val json: JsonElement
		abstract val file: IStarboundFile?
		abstract val registry: Registry<T>
		@Deprecated("Careful! This might be confused with 'isMeta' of some entries (such as tiles, in such case use entry.isMeta)")
		abstract val isBuiltin: Boolean
		abstract val ref: Ref<T>

		final override fun get(): T {
			return value
		}

		val jsonObject: JsonObject
			get() = json as JsonObject
	}

	private inner class Impl(override val key: String, override var value: T, override var id: Int? = null) : Entry<T>() {
		override var json: JsonElement = JsonNull.INSTANCE
		override var file: IStarboundFile? = null
		override var isBuiltin: Boolean = false

		override val ref: Ref<T> by lazy { ref(key) }

		override fun equals(other: Any?): Boolean {
			return this === other
		}

		private val hash: Int

		init {
			var x = key.hashCode()
			// avalanche bits using murmur3 hash
			x = x xor (x ushr 16)
			x *= -0x7a143595
			x = x xor (x ushr 13)
			x *= -0x3d4d51cb
			x = x xor (x ushr 16)
			hash = x
		}

		override fun hashCode(): Int {
			return hash
		}

		override fun toString(): String {
			return "Entry of $name at $key/${id ?: "-"}"
		}

		override val registry: Registry<T>
			get() = this@Registry
	}

	private inner class RefImpl(override val key: Either<String, Int>) : Ref<T>() {
		override var entry: Entry<T>? = null
		var references = 0

		override fun equals(other: Any?): Boolean {
			return this === other || other is Registry<*>.RefImpl && other.key == key && other.registry == registry
		}

		private val hash: Int

		init {
			var x = key.hashCode()
			// avalanche bits using murmur3 hash
			x = x xor (x ushr 16)
			x *= -0x7a143595
			x = x xor (x ushr 13)
			x *= -0x3d4d51cb
			x = x xor (x ushr 16)
			hash = x
		}

		override fun hashCode(): Int {
			return hash
		}

		override fun toString(): String {
			return "Ref of $name at $key/${if (entry != null) "bound" else "missing"}"
		}

		override val registry: Registry<T>
			get() = this@Registry
	}

	operator fun get(index: String): Entry<T>? {
		val result = lock.read { keysInternal[index] }

		if (result == null && hasBeenValidated) {
			val hasher = XXHash64()
			hasher.update(index.toByteArray())
			val missIndex = hasher.digestAsLong()

			lock.write {
				if (loggedMisses.add(missIndex)) {
					LOGGER.warn("No such $name: ${index.limit()}")
				}
			}
		}

		return result
	}

	operator fun get(index: Int): Entry<T>?  {
		val result = lock.read { idsInternal[index] }

		if (result == null && hasBeenValidated) {
			val hasher = XXHash64()
			hasher.update(index.toByte())
			hasher.update((index ushr 8).toByte())
			hasher.update((index ushr 16).toByte())
			hasher.update((index ushr 24).toByte())
			val missIndex = hasher.digestAsLong()

			lock.write {
				if (loggedMisses.add(missIndex)) {
					LOGGER.warn("No such $name: ID $index")
				}
			}
		}

		return result
	}

	fun getOrThrow(index: String): Entry<T> {
		return get(index) ?: throw NoSuchElementException("No such $name: $index")
	}

	fun getOrThrow(index: Int): Entry<T> {
		return get(index) ?: throw NoSuchElementException("No such $name: $index")
	}

	/**
	 * To be used inside network readers, so clients sending invalid data will get kicked
	 */
	fun refOrThrow(index: String): Ref<T> = get(index)?.ref ?: throw NoSuchElementException("No such $name: ${index.limit()}")

	fun ref(index: String): Ref<T> = lock.write {
		keyRefs.computeIfAbsent(index, Object2ObjectFunction {
			val ref = RefImpl(Either.left(it as String))
			ref.entry = keysInternal[it]

			if (hasBeenValidated && ref.entry == null) {
				val hasher = XXHash64()
				hasher.update(index.toByteArray())

				if (loggedMisses.add(hasher.digestAsLong())) {
					LOGGER.warn("No such $name: ${it.limit()}")
				}
			}

			ref
		}).also {
			it.references++
		}
	}

	/**
	 * To be used inside network readers, so clients sending invalid data will get kicked
	 */
	fun refOrThrow(index: Int): Ref<T> = get(index)?.ref ?: throw NoSuchElementException("No such $name: ID $index")

	fun ref(index: Int): Ref<T> = lock.write {
		idRefs.computeIfAbsent(index, Int2ObjectFunction {
			val ref = RefImpl(Either.right(it))
			ref.entry = idsInternal[it]

			if (hasBeenValidated && ref.entry == null) {
				val hasher = XXHash64()
				hasher.update(index.toByte())
				hasher.update((index ushr 8).toByte())
				hasher.update((index ushr 16).toByte())
				hasher.update((index ushr 24).toByte())

				if (loggedMisses.add(hasher.digestAsLong())) {
					LOGGER.warn("No such $name: ID $it")
				}
			}

			ref
		}).also {
			it.references++
		}
	}

	operator fun contains(index: String) = lock.read { index in keysInternal }
	operator fun contains(index: Int) = lock.read { index in idsInternal }

	fun validate(): Boolean {
		hasBeenValidated = true

		lock.read {
			var valid = true

			keyRefs.values.forEach {
				if (!it.isPresent && it.key.left() != "") {
					LOGGER.warn("Registry '$name' reference at '${it.key.left()}' is not bound to value, expect problems (referenced ${it.references} times)")
					valid = false
				}
			}

			idRefs.values.forEach {
				if (!it.isPresent) {
					LOGGER.warn("Registry '$name' reference with ID '${it.key.right()}' is not bound to value, expect problems (referenced ${it.references} times)")
					valid = false
				}
			}

			return valid
		}
	}

	fun add(
		key: String,
		value: T,
		json: JsonElement = JsonNull.INSTANCE,
		file: IStarboundFile? = null,
		id: KOptional<Int?> = KOptional(),
		isBuiltin: Boolean = false
	): Entry<T> {
		require(key != "") { "Adding $name with empty name (empty name is reserved)" }

		lock.write {
			if (key in keysInternal) {
				LOGGER.warn("Overwriting $name at '$key' (new def originate from $file; old def originate from ${keysInternal[key]?.file ?: "<code>"})")
			}

			id.ifPresent { id ->
				if (id != null && id in idsInternal) {
					LOGGER.warn("Overwriting $name with ID '$id' (new def originate from $file; old def originate from ${idsInternal[id]?.file ?: "<code>"})")
				}
			}

			val entry = keysInternal.computeIfAbsent(key, Object2ObjectFunction { Impl(key, value) })
			check(!entry.isBuiltin || isBuiltin) { "Trying to redefine builtin entry at $key" }

			id.ifPresent {
				entry.id?.let {
					idsInternal.remove(it)
					idRefs[it]?.entry = null
				}

				entry.id = it
			}

			entry.value = value

			if (storeJson)
				entry.json = json

			entry.file = file
			entry.isBuiltin = isBuiltin

			keyRefs[key]?.entry = entry

			id.ifPresent { id ->
				if (id != null) {
					idRefs[id]?.entry = entry
					idsInternal[id] = entry
				}
			}

			return entry
		}
	}

	val emptyRef = ref("")

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