package ru.dbotthepony.mc.otm.network.syncher

import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap
import it.unimi.dsi.fastutil.ints.IntArraySet
import it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet
import net.minecraft.network.RegistryFriendlyByteBuf
import ru.dbotthepony.kommons.util.KOptional
import ru.dbotthepony.mc.otm.core.IDAllocator
import java.io.Closeable
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.atomic.AtomicBoolean

/**
 * Syncher group/set, which deals with synchables of only one type, and which are created and removed
 * on remote (e.g. server adding and removing synchables at will), which makes it distinct from
 * [SynchableGroup], in which attached synchables are created/removed manually on both sides.
 */
class DynamicSynchableGroup<T : ISynchable>(
	/**
	 * Constructs new [T] instance locally, when remote created one.
	 * Data written by [writer] must be read here, if there is any.
	 */
	private val reader: RegistryFriendlyByteBuf.() -> T,

	/**
	 * Allows to write additional data to network stream during
	 * first-time networking of [T] to remote
	 */
	private val writer: T.(RegistryFriendlyByteBuf) -> Unit = {}
) : ISynchable, MutableSet<T> {
	constructor(factory: () -> T) : this({ factory() }, {})

	private inner class RemoteState(val listener: Runnable) : IRemoteState {
		private inner class RemoteSlot(val slot: Slot<T>, fromConstructor: Boolean) : Runnable, Closeable {
			val remoteState = slot.synchable.createRemoteState(this)
			val isDirty = AtomicBoolean(true)
			var isRemoved = false
				private set

			init {
				if (!fromConstructor)
					removals.remove(slot.id)

				firstTime.add(this)
				dirty.add(this)

				if (firstTime.size != 1 && !fromConstructor) {
					listener.run()
				}
			}

			override fun run() {
				if (isDirty.compareAndSet(false, true)) {
					dirty.add(this)
					listener.run()
				}
			}

			fun markClean() {
				isDirty.set(false)
			}

			override fun close() {
				isRemoved = true
				isDirty.set(true)
				remoteState.close()
			}
		}

		private val remotes = HashMap<Slot<T>, RemoteSlot>()
		private val dirty = ConcurrentLinkedQueue<RemoteSlot>()
		private val firstTime = ArrayList<RemoteSlot>()
		private val removals = IntArraySet()
		private var sendClear = false

		init {
			value2slot.values.forEach {
				remotes[it] = RemoteSlot(it, true)
			}

			remoteStates.add(this)
		}

		fun add(element: Slot<T>) {
			remotes[element] = RemoteSlot(element, false)
			listener.run()
		}

		fun remove(element: Slot<T>) {
			val removed = remotes.remove(element)
			checkNotNull(removed)
			check(removed.slot == element)
			removed.close()
			listener.run()
		}

		override fun write(stream: RegistryFriendlyByteBuf) {
			if (sendClear) {
				sendClear = false
				stream.writeByte(CLEAR)
			}

			run {
				val itr = removals.iterator()

				while (itr.hasNext()) {
					stream.writeByte(REMOVE_ENTRY)
					stream.writeVarInt(itr.nextInt())
				}

				removals.clear()
			}

			firstTime.forEach {
				stream.writeByte(ADD_ENTRY)
				stream.writeVarInt(it.slot.id)
				writer(it.slot.synchable, stream)
			}

			firstTime.clear()

			val seen = ReferenceOpenHashSet<RemoteSlot>()
			val slots = ArrayList<RemoteSlot>()
			var next = dirty.poll()

			while (next != null) {
				if (!next.isRemoved) {
					val status = seen.add(next)
					next.markClean()

					if (status) {
						slots.add(next)
					}
				}

				next = dirty.poll()
			}

			slots.forEach {
				stream.writeByte(SYNC_ENTRY)
				stream.writeVarInt(it.slot.id)
				it.remoteState.write(stream)
			}

			stream.writeByte(END)
		}

		override fun invalidate() {
			firstTime.clear()
			dirty.clear()
			removals.clear() // is it necessary?

			remotes.values.forEach {
				firstTime.add(it)
				dirty.add(it)
			}
		}

		override fun close() {
			remotes.values.forEach { it.close() }

			firstTime.clear()
			dirty.clear()
			removals.clear()
			remotes.clear()

			remoteStates.remove(this)
		}

		fun clear() {
			sendClear = true
			firstTime.clear()
			dirty.clear()
			removals.clear()
			remotes.values.forEach { it.close() }
			remotes.clear()
		}
	}

	private data class Slot<T : Any>(val synchable: T, val id: Int)

	private val remoteStates = ArrayList<RemoteState>()
	private val value2slot = HashMap<T, Slot<T>>()
	private val id2slot = Int2ObjectOpenHashMap<Slot<T>>()
	private val idAllocator = IDAllocator()

	override val hasRemotes: Boolean
		get() = remoteStates.isNotEmpty()

	override var isRemote: Boolean = false
		private set

	override fun read(stream: RegistryFriendlyByteBuf) {
		isRemote = true
		var action = stream.readByte().toInt()

		while (action != END) {
			when (action) {
				REMOVE_ENTRY -> {
					val id = stream.readVarInt()
					val slot = checkNotNull(id2slot.remove(id)) { "No such slot with ID: $id" }
					check(value2slot.remove(slot.synchable) == value2slot)
					remoteStates.forEach { it.remove(slot) }
				}

				ADD_ENTRY -> {
					val id = stream.readVarInt()

					// it is expected that during very frequent updates entry gets removed and new entry takes up their ID,
					// so we allow implicit removal if ID is allocated by some other entry before old entry could have a chance
					// to broadcast that it have been removed

					if (id2slot.containsKey(id)) {
						val slot = id2slot.remove(id)!!
						check(value2slot.remove(slot.synchable) == value2slot)
						remoteStates.forEach { it.remove(slot) }
					}

					val value = reader(stream)
					val slot = Slot(value, id)
					value2slot[value] = slot
					id2slot[id] = slot
				}

				SYNC_ENTRY -> {
					val id = stream.readVarInt()
					val slot = checkNotNull(id2slot.get(id)) { "No such slot with ID: $id" }
					slot.synchable.read(stream)
				}

				CLEAR -> {
					value2slot.clear()
					id2slot.clear()
				}

				else -> throw IllegalArgumentException("Unknown network action: $action")
			}

			action = stream.readByte().toInt()
		}
	}

	override fun createRemoteState(listener: Runnable): IRemoteState {
		return RemoteState(listener)
	}

	override fun add(element: T): Boolean {
		if (element in value2slot) {
			return false
		}

		val slot = Slot(element, idAllocator.allocate())
		value2slot[element] = slot
		id2slot[slot.id] = slot
		remoteStates.forEach { it.add(slot) }
		return true
	}

	override val size: Int
		get() = value2slot.size

	override fun addAll(elements: Collection<T>): Boolean {
		var any = false
		elements.forEach { any = add(it) || any }
		return any
	}

	override fun clear() {
		if (value2slot.isNotEmpty()) {
			value2slot.clear()
			id2slot.clear()

			remoteStates.toTypedArray().forEach { it.clear() }
		}
	}

	override fun contains(element: T): Boolean {
		return element in value2slot
	}

	override fun containsAll(elements: Collection<T>): Boolean {
		return value2slot.keys.containsAll(elements)
	}

	override fun isEmpty(): Boolean {
		return value2slot.isEmpty()
	}

	override fun iterator(): MutableIterator<T> {
		return object : MutableIterator<T> {
			private val parent = value2slot.values.iterator()
			private var last: KOptional<Slot<T>> = KOptional()

			override fun hasNext(): Boolean {
				return parent.hasNext()
			}

			override fun next(): T {
				val slot = parent.next()
				last = KOptional(slot)
				return slot.synchable
			}

			override fun remove() {
				parent.remove()
				val slot = last.value
				checkNotNull(id2slot.remove(slot.id))
				remoteStates.forEach { it.remove(slot) }
				last = KOptional()
			}
		}
	}

	override fun remove(element: T): Boolean {
		if (element !in value2slot) {
			return false
		}

		val slot = value2slot.remove(element)!!
		checkNotNull(id2slot.remove(slot.id))
		remoteStates.forEach { it.remove(slot) }
		return true
	}

	override fun removeAll(elements: Collection<T>): Boolean {
		var any = false
		elements.forEach { any = remove(it) || any }
		return any
	}

	override fun retainAll(elements: Collection<T>): Boolean {
		val itr = iterator()
		var any = false

		for (element in itr) {
			if (element !in elements) {
				itr.remove()
				any = true
			}
		}

		return any
	}

	companion object {
		private const val END = 0
		private const val ADD_ENTRY = 1
		private const val SYNC_ENTRY = 2
		private const val REMOVE_ENTRY = 3
		private const val CLEAR = 4
	}
}
