package ru.dbotthepony.kstarbound.network.syncher

import it.unimi.dsi.fastutil.ints.IntAVLTreeSet
import it.unimi.dsi.fastutil.io.FastByteArrayInputStream
import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
import ru.dbotthepony.kommons.io.readByteArray
import ru.dbotthepony.kommons.io.readVarInt
import ru.dbotthepony.kommons.io.writeByteArray
import ru.dbotthepony.kommons.io.writeVarInt
import ru.dbotthepony.kstarbound.collect.IdMap
import java.io.DataInputStream
import java.io.DataOutputStream
import java.util.Collections
import java.util.function.LongSupplier

/**
 * A dynamic group of NetElements that manages creation and destruction of
 * individual elements, that is itself a NetElement.  Element changes are not
 * delayed by the interpolation delay, they will always happen immediately, but
 * this does not inhibit the Elements themselves from handling their own delta
 * update delays normally.
 */
class NetworkedDynamicGroup<T : Any>(private val factory: () -> T, private val accessor: (T) -> NetworkedElement, private val maxBacklogSize: Int = 100) : NetworkedElement() {
	private val elementsInternal = IdMap<T>(min = 1)
	val elements: Map<Int, T> = Collections.unmodifiableMap(elementsInternal)
	private val backlog = ArrayDeque<Pair<Long, Entry>>()
	private var isInterpolating = false
	private var extrapolation = 0.0

	private enum class Action {
		CLEAR, REMOVE, ADD;
	}

	override fun toString(): String {
		return "NetworkedDynamicGroup[keys = ${elementsInternal.size}]"
	}

	// Storing data seems to be redundant at first, but this actually makes sense
	// since networked elements might have interpolation enabled
	//
	// If we would store element itself in Entry, and write it's initial state
	// upon networking to remote, we would skip all interpolation queues, since
	// IF element has interpolation enabled, writeInitial will ALWAYS write latest-possible
	// data (data from end of interpolation queue), and this is not something desired
	// (because writes of this data happens on both writeInitial and writeDelta inside this class)
	// Major downside of this is keeping byte arrays around
	private class Entry(
		val action: Action,
		val id: Int,
		val dataNative: ByteArray?,
		val dataLegacy: ByteArray?,
	) {
		constructor(id: Int, element: NetworkedElement) : this(
			Action.ADD,
			id,
			FastByteArrayOutputStream().let {
				element.writeInitial(DataOutputStream(it), false)
				it.array.copyOf(it.length)
			},
			FastByteArrayOutputStream().let {
				element.writeInitial(DataOutputStream(it), true)
				it.array.copyOf(it.length)
			})

		constructor(id: Int) : this(Action.REMOVE, id, null, null)
	}

	private fun setupElement(element: NetworkedElement) {
		val versionCounter = versionCounter

		if (versionCounter != null)
			element.specifyVersioner(versionCounter)

		if (isInterpolating)
			element.enableInterpolation(extrapolation)
		else
			element.disableInterpolation()
	}

	fun add(element: T): Int {
		check(element !in elementsInternal.values) { "Already containing $element" }
		setupElement(accessor(element))
		val id = elementsInternal.add(element)
		backlog.add(currentVersion() to Entry(id, accessor(element)))
		purgeBacklog()
		return id
	}

	fun remove(index: Int): Boolean {
		return elementsInternal.remove(index) != null
	}

	private fun purgeBacklog() {
		while (backlog.size >= maxBacklogSize) {
			backlog.removeFirst()
		}
	}

	override fun hasChangedSince(version: Long): Boolean {
		return version == 0L || backlog.any { it.first >= version } || elementsInternal.values.any { accessor(it).hasChangedSince(version) }
	}

	override fun specifyVersioner(versionCounter: LongSupplier) {
		super.specifyVersioner(versionCounter)
		backlog.clear()
		backlog.add(currentVersion() to clearAction)

		for ((id, element) in elementsInternal) {
			accessor(element).specifyVersioner(versionCounter)
			backlog.add(currentVersion() to Entry(id, accessor(element)))
		}

		purgeBacklog()
	}

	override fun readInitial(data: DataInputStream, isLegacy: Boolean) {
		val size = data.readVarInt()

		backlog.clear()
		elementsInternal.clear()

		backlog.add(currentVersion() to clearAction)

		for (i in 0 until size) {
			val index = data.readVarInt()

			val stream = if (isLegacy) {
				// sigh
				DataInputStream(FastByteArrayInputStream(data.readByteArray()))
			} else {
				data
			}

			val element = factory()
			setupElement(accessor(element))
			elementsInternal[index] = element
			accessor(element).readInitial(stream, isLegacy)
			backlog.add(currentVersion() to Entry(index, accessor(element)))
		}
	}

	override fun writeInitial(data: DataOutputStream, isLegacy: Boolean) {
		data.writeVarInt(elementsInternal.size)

		for ((id, element) in elementsInternal) {
			data.writeVarInt(id)

			if (isLegacy) {
				// sigh
				val wrap = FastByteArrayOutputStream()
				accessor(element).writeInitial(DataOutputStream(wrap), true)
				data.writeByteArray(wrap.array, 0, wrap.length)
			} else {
				accessor(element).writeInitial(data, false)
			}
		}
	}

	override fun readDelta(data: DataInputStream, interpolationDelay: Double, isLegacy: Boolean) {
		if (data.readBoolean()) {
			readInitial(data, isLegacy)
		} else {
			val received = IntAVLTreeSet()

			while (true) {
				when (val id = data.readVarInt()) {
					0 -> break

					1 -> {
						when (Action.entries[data.readUnsignedByte()]) {
							Action.CLEAR -> {
								elementsInternal.clear()
								backlog.add(currentVersion() to clearAction)
							}

							Action.REMOVE -> {
								// inconsistent usage of VarInt and int32_t when networking is driving me insane
								val id = if (isLegacy) data.readInt() else data.readVarInt()
								backlog.add(currentVersion() to Entry(id))
								elementsInternal.remove(id)
							}

							Action.ADD -> {
								// inconsistent usage of VarInt and int32_t when networking is driving me insane
								val id = if (isLegacy) data.readInt() else data.readVarInt()
								val element = factory()
								setupElement(accessor(element))
								check(id !in elementsInternal) { "Already has networked element $id" }
								elementsInternal[id] = element

								if (isLegacy) {
									accessor(element).readInitial(DataInputStream(FastByteArrayInputStream(data.readByteArray())), true)
								} else {
									accessor(element).readInitial(data, false)
								}

								backlog.add(currentVersion() to Entry(id, accessor(element)))
							}
						}
					}

					else -> {
						val element = elementsInternal[id - 1] ?: throw NoSuchElementException("Unknown network element with ID ${id - 1}, net state is corrupt!")
						accessor(element).readDelta(data, interpolationDelay, isLegacy)

						if (isInterpolating) {
							received.add(id)
						}
					}
				}
			}

			if (isInterpolating) {
				for ((id, element) in elementsInternal) {
					if (!received.contains(id)) {
						accessor(element).readBlankDelta(interpolationDelay)
					}
				}
			}
		}
	}

	override fun writeDelta(data: DataOutputStream, remoteVersion: Long, isLegacy: Boolean) {
		// backlog is guaranteed to be not empty once first networked element is added
		if (backlog.isNotEmpty() && remoteVersion < backlog.first().first) {
			// we fell way behind, write full state
			data.writeBoolean(true)
			writeInitial(data, isLegacy)
			return
		}

		data.writeBoolean(false)

		for ((version, entry) in backlog) {
			if (version >= remoteVersion) {
				data.writeByte(1)
				data.writeByte(entry.action.ordinal)

				when (entry.action) {
					Action.CLEAR -> {}
					Action.REMOVE -> if (isLegacy) data.writeInt(entry.id) else data.writeVarInt(entry.id)

					Action.ADD -> {
						if (isLegacy) {
							data.writeInt(entry.id)
							data.writeByteArray(entry.dataLegacy!!)
						} else {
							data.writeVarInt(entry.id)
							data.write(entry.dataNative!!)
						}
					}
				}
			}
		}

		for ((id, element) in elementsInternal) {
			if (accessor(element).hasChangedSince(remoteVersion)) {
				data.writeVarInt(id + 1)
				accessor(element).writeDelta(data, remoteVersion, isLegacy)
			}
		}

		data.writeByte(0)
	}

	override fun readBlankDelta(interpolationDelay: Double) {
		elementsInternal.values.forEach { accessor(it).readBlankDelta(interpolationDelay) }
	}

	override fun enableInterpolation(extrapolation: Double) {
		isInterpolating = true
		this.extrapolation = extrapolation
		elementsInternal.values.forEach { accessor(it).enableInterpolation(extrapolation) }
	}

	override fun disableInterpolation() {
		isInterpolating = false
		elementsInternal.values.forEach { accessor(it).disableInterpolation() }
	}

	override fun tickInterpolation(delta: Double) {
		elementsInternal.values.forEach { accessor(it).tickInterpolation(delta) }
	}

	companion object {
		private val clearAction = Entry(Action.CLEAR, 0, null, null)
	}
}
