package ru.dbotthepony.mc.otm.block.entity

import it.unimi.dsi.fastutil.booleans.BooleanConsumer
import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap
import it.unimi.dsi.fastutil.longs.Long2ObjectFunction
import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap
import it.unimi.dsi.fastutil.objects.Object2ObjectArrayMap
import it.unimi.dsi.fastutil.objects.ObjectArraySet
import it.unimi.dsi.fastutil.objects.Reference2IntArrayMap
import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
import net.minecraft.core.BlockPos
import net.minecraft.core.Direction
import net.minecraft.core.HolderLookup.Provider
import net.minecraft.core.SectionPos
import net.minecraft.core.Vec3i
import net.minecraft.nbt.CompoundTag
import net.minecraft.server.level.ServerLevel
import net.minecraft.server.level.ServerPlayer
import net.minecraft.world.Container
import net.minecraft.world.entity.player.Player
import net.minecraft.world.level.ChunkPos
import net.minecraft.world.level.Level
import net.minecraft.world.level.block.Block
import net.minecraft.world.level.block.entity.BlockEntity
import net.minecraft.world.level.block.entity.BlockEntityType
import net.minecraft.world.level.block.entity.RandomizableContainerBlockEntity
import net.minecraft.world.level.block.state.BlockState
import net.minecraft.world.level.chunk.LevelChunk
import net.minecraft.world.phys.Vec3
import net.neoforged.neoforge.capabilities.BlockCapability
import net.neoforged.neoforge.capabilities.BlockCapabilityCache
import net.neoforged.neoforge.capabilities.Capabilities
import net.neoforged.neoforge.event.entity.player.PlayerEvent
import net.neoforged.neoforge.event.level.ChunkWatchEvent
import net.neoforged.neoforge.event.level.LevelEvent
import net.neoforged.neoforge.event.server.ServerStoppingEvent
import net.neoforged.neoforge.event.tick.LevelTickEvent
import net.neoforged.neoforge.network.PacketDistributor
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.kommons.util.Listenable
import ru.dbotthepony.mc.otm.SERVER_IS_LIVE
import ru.dbotthepony.mc.otm.block.INeighbourChangeListener
import ru.dbotthepony.mc.otm.block.RotatableMatteryBlock
import ru.dbotthepony.mc.otm.capability.MatteryCapability
import ru.dbotthepony.mc.otm.capability.energy.IMatteryEnergyStorage
import ru.dbotthepony.mc.otm.core.collect.WeakHashSet
import ru.dbotthepony.mc.otm.core.get
import ru.dbotthepony.mc.otm.core.math.BlockRotation
import ru.dbotthepony.mc.otm.core.math.RelativeSide
import ru.dbotthepony.mc.otm.core.math.plus
import ru.dbotthepony.mc.otm.core.util.IntCounter
import ru.dbotthepony.mc.otm.core.util.Savetables
import ru.dbotthepony.mc.otm.core.util.TickList
import ru.dbotthepony.mc.otm.core.util.countingLazy
import ru.dbotthepony.mc.otm.network.BlockEntitySyncPacket
import ru.dbotthepony.mc.otm.network.syncher.SynchableGroup
import ru.dbotthepony.mc.otm.once
import ru.dbotthepony.mc.otm.onceServer
import ru.dbotthepony.mc.otm.registry.game.MBlocks
import ru.dbotthepony.mc.otm.sometimeServer
import java.lang.ref.WeakReference
import java.util.*
import java.util.function.BooleanSupplier
import java.util.function.Consumer
import java.util.function.Predicate
import java.util.function.Supplier
import java.util.stream.Stream
import kotlin.collections.ArrayList

/**
 * Absolute barebone (lol) block entity class in Overdrive that Matters, providing bare minimum (lulmao, minecraft engine) functionality
 */
abstract class MatteryBlockEntity(p_155228_: BlockEntityType<*>, p_155229_: BlockPos, p_155230_: BlockState) : BlockEntity(p_155228_, p_155229_, p_155230_), INeighbourChangeListener {
	private val sidelessCaps = Reference2ObjectOpenHashMap<BlockCapability<*, *>, Any>()
	private val sidedCaps = Array(RelativeSide.entries.size) {
		Reference2ObjectOpenHashMap<BlockCapability<*, *>, ControllableCapability<*>>()
	}

	protected val tickList = TickList()
	protected val blockStateChangesCounter = IntCounter()
	protected val dirtyListeners = Listenable.Impl<Unit>()
	private val waitForServerLevel = ArrayList<() -> Unit>()

	private val _droppableContainers = ObjectArraySet<Container>()
	private val _neighbourChangeListeners = ObjectArraySet<INeighbourChangeListener>()
	val droppableContainers: Set<Container> = Collections.unmodifiableSet(_droppableContainers)
	val neighbourChangeListeners: Set<INeighbourChangeListener> = Collections.unmodifiableSet(_neighbourChangeListeners)

	/**
	 * Shared savetables, written both to level storage and to item tag
	 */
	protected val savetables = Savetables()

	/**
	 * Level-only savetables, written only to level storage
	 */
	protected val savetablesLevel = Savetables()

	/**
	 * Savetables for things configured from user input
	 */
	protected val savetablesConfig = Savetables()

	/**
	 * "shortcut" for getting [BlockRotation]
	 *
	 * if block has no rotation, returns [BlockRotation.NORTH]
	 */
	open val blockRotation: BlockRotation by countingLazy(blockStateChangesCounter) {
		blockState[(blockState.block as? RotatableMatteryBlock ?: return@countingLazy BlockRotation.NORTH).rotationProperty]
	}

	protected fun addDroppableContainer(container: Container) {
		_droppableContainers.add(container)
	}

	protected fun addNeighbourListener(listener: INeighbourChangeListener) {
		if (listener === this)
			throw IllegalArgumentException("are you drunk")

		_neighbourChangeListeners.add(listener)
	}

	open fun beforeDroppingItems(oldBlockState: BlockState, level: Level, blockPos: BlockPos, newBlockState: BlockState, movedByPiston: Boolean) {}
	open fun beforeDestroyedByPlayer(level: Level, blockPos: BlockPos, blockState: BlockState, player: Player) {}

	open fun tick() {
		tickList.tick()
	}

	interface CapabilityControl : BooleanSupplier, BooleanConsumer {
		var isEnabled: Boolean

		fun doDispatchInvalidation()
		fun dontDispatchInvalidation()

		override fun getAsBoolean(): Boolean {
			return isEnabled
		}

		override fun accept(t: Boolean) {
			isEnabled = t
		}
	}

	private inner class ControllableCapability<T : Any>(val capability: T) : CapabilityControl {
		override var isEnabled: Boolean = true
			set(value) {
				if (value != field) {
					field = value
					if (dispatchInvalidation) level?.invalidateCapabilities(blockPos)
				}
			}

		private var dispatchInvalidation = true

		override fun doDispatchInvalidation() {
			dispatchInvalidation = true
		}

		override fun dontDispatchInvalidation() {
			dispatchInvalidation = false
		}

		fun getOrNull(): T? {
			return if (isEnabled) capability else null
		}
	}

	private inner class ControllableCapabilitySet(val capabilities: List<CapabilityControl>) : CapabilityControl {
		override var isEnabled: Boolean
			get() = capabilities.all { it.isEnabled }
			set(value) {
				capabilities.forEach { it.dontDispatchInvalidation() }

				try {
					capabilities.forEach { it.isEnabled = value }
				} finally {
					capabilities.forEach { it.doDispatchInvalidation() }
					level?.invalidateCapabilities(blockPos)
				}
			}

		override fun doDispatchInvalidation() {
			capabilities.forEach { it.doDispatchInvalidation() }
		}

		override fun dontDispatchInvalidation() {
			capabilities.forEach { it.dontDispatchInvalidation() }
		}
	}

	/**
	 * exposes capability when no side is specified
	 */
	protected fun <T : Any> exposeSideless(capability: BlockCapability<T, *>, value: T) {
		check(!sidelessCaps.containsKey(capability)) { "Already has globally exposed $capability!" }
		MBlocks.ensureCapabilityIsKnown(capability)
		sidelessCaps[capability] = value
		setChanged()
		level?.invalidateCapabilities(blockPos)
	}

	protected fun <T : Any> exposeSided(side: RelativeSide, capability: BlockCapability<T, *>, value: T): CapabilityControl {
		val map = sidedCaps[side.ordinal]
		check(!map.containsKey(capability)) { "Already has exposed $capability on $side!" }
		MBlocks.ensureCapabilityIsKnown(capability)
		val wrapper = ControllableCapability(value)
		map[capability] = wrapper
		setChanged()
		level?.invalidateCapabilities(blockPos)
		return wrapper
	}

	/**
	 * Exposes capability unconditionally, on all sides and sideless
	 */
	protected fun <T : Any> exposeGlobally(capability: BlockCapability<T, *>, value: T, predicate: Predicate<RelativeSide> = Predicate { true }): CapabilityControl {
		exposeSideless(capability, value)

		return ControllableCapabilitySet(RelativeSide.entries.stream().filter(predicate).map { exposeSided(it, capability, value) }.toList())
	}

	protected fun exposeEnergySideless(value: IMatteryEnergyStorage) {
		exposeSideless(Capabilities.EnergyStorage.BLOCK, value)
		exposeSideless(MatteryCapability.BLOCK_ENERGY, value)
	}

	protected fun exposeEnergyGlobally(value: IMatteryEnergyStorage): CapabilityControl {
		return ControllableCapabilitySet(
			listOf(
				exposeGlobally(Capabilities.EnergyStorage.BLOCK, value),
				exposeGlobally(MatteryCapability.BLOCK_ENERGY, value)
			)
		)
	}

	protected fun exposeEnergySided(side: RelativeSide, value: IMatteryEnergyStorage): CapabilityControl {
		return ControllableCapabilitySet(
			listOf(
				exposeSided(side, Capabilities.EnergyStorage.BLOCK, value),
				exposeSided(side, MatteryCapability.BLOCK_ENERGY, value)
			)
		)
	}

	protected fun waitForServerLevel(lambda: () -> Unit) {
		if (level is ServerLevel) {
			if (executeWaitForServerLevel)
				lambda.invoke()
			else
				waitForServerLevel.add(lambda)
		} else if (level == null) {
			waitForServerLevel.add(lambda)
		}
	}

	fun <T : Any> getCapability(cap: BlockCapability<T, *>, side: Direction?): T? {
		if (side != null) {
			return sidedCaps[blockRotation.dir2Side(side).ordinal][cap]?.getOrNull() as T?
		}

		return sidelessCaps[cap] as T?
	}

	final override fun saveAdditional(nbt: CompoundTag, registry: Provider) {
		super.saveAdditional(nbt, registry)
		saveShared(nbt, registry)
		saveLevel(nbt, registry)
	}

	/**
	 * Saved both to item dropped, and to level storage
	 */
	open fun saveShared(nbt: CompoundTag, registry: Provider) {
		savetables.serializeNBT(nbt, registry)
		saveConfiguration(nbt, registry)
	}

	/**
	 * Only saved to level storage, discarded when dropped as item
	 */
	open fun saveLevel(nbt: CompoundTag, registry: Provider) {
		savetablesLevel.serializeNBT(nbt, registry)
	}

	override fun loadAdditional(nbt: CompoundTag, registry: Provider) {
		super.loadAdditional(nbt, registry)
		savetables.deserializeNBT(registry, nbt)
		savetablesLevel.deserializeNBT(registry, nbt)
		loadConfiguration(nbt, registry)
	}

	fun saveConfiguration(nbt: CompoundTag, registry: Provider) {
		savetablesConfig.serializeNBT(nbt, registry)
	}

	fun loadConfiguration(nbt: CompoundTag, registry: Provider) {
		savetablesConfig.deserializeNBT(registry, nbt)
	}

	@Suppress("OVERRIDE_DEPRECATION")
	override fun setBlockState(pBlockState: BlockState) {
		val old = blockRotation
		blockStateChangesCounter.increment()
		@Suppress("DEPRECATION")
		super.setBlockState(pBlockState)
		val new = blockRotation

		if (old != new) {
			level?.invalidateCapabilities(blockPos)
			capabilityCaches.forEach { it.rebuildCache() }
		}
	}

	override fun neighborChanged(state: BlockState, level: Level, pos: BlockPos, neighbour: Block, neighbourPos: BlockPos, movedByPiston: Boolean) {
		_neighbourChangeListeners.forEach { it.neighborChanged(state, level, pos, neighbour, neighbourPos, movedByPiston) }
	}

	override fun setChanged() {
		super.setChanged()

		dirtyListeners.accept(Unit)
	}

	// Just to mark chunk unsaved
	open fun markDirtyFast() {
		val level = level

		if (level is ServerLevel) {
			level.chunkSource.getChunkNow(
				SectionPos.blockToSectionCoord(blockPos.x),
				SectionPos.blockToSectionCoord(blockPos.z))?.isUnsaved = true
		}

		dirtyListeners.accept(Unit)
	}

	private val capabilityCaches = ArrayList<CapabilityCache<*>>()

	inner class CapabilityCache<T : Any>(val side: RelativeSide, val capability: BlockCapability<T, in Direction?>) : Supplier<T?>, Listenable<T?> {
		private var currentVersion = 0
		private var cache: BlockCapabilityCache<T, in Direction?>? = null
		private val listeners = Listenable.Impl<T?>()

		override fun addListener(listener: Consumer<T?>): Listenable.L {
			return listeners.addListener(listener)
		}

		init {
			capabilityCaches.add(this)
			waitForServerLevel { rebuildCache() }
		}

		override fun get(): T? {
			return cache?.capability
		}

		val isPresent: Boolean
			get() = cache?.capability != null

		val isEmpty: Boolean
			get() = cache?.capability == null

		fun rebuildCache() {
			if (!SERVER_IS_LIVE) return
			val level = level as? ServerLevel

			if (level == null) {
				cache = null
				return
			}

			val creationVersion = ++currentVersion

			val direction = blockRotation.side2Dir(side)

			cache = BlockCapabilityCache.create(
				capability, level, blockPos + direction.normal, direction.opposite,
				{ !isRemoved || creationVersion != currentVersion },
				// IllegalStateException("Do not call getCapability on an invalid cache or from the invalidation listener!")
				// what a shame.
				{ if (SERVER_IS_LIVE) onceServer { if (!isRemoved && creationVersion == currentVersion) listeners.accept(cache?.capability) } })

			onceServer {
				if (!isRemoved && creationVersion == currentVersion) listeners.accept(cache?.capability)
			}
		}
	}

	val syncher = SynchableGroup()
	private val synchers = Object2ObjectArrayMap<ServerPlayer, SynchableGroup.Remote>()
	private var executeWaitForServerLevel = false

	override fun setLevel(level: Level) {
		val old = this.level
		super.setLevel(level)
		unsubscribe()
		_subCache = null

		if (!level.isClientSide) {
			subscribe()

			level.once {
				if (isRemoved) return@once
				executeWaitForServerLevel = true
				waitForServerLevel.forEach { it.invoke() }
				waitForServerLevel.clear()
			}
		} else {
			executeWaitForServerLevel = true
			waitForServerLevel.clear()
			BlockEntitySyncPacket.applyBacklog(this)
		}
	}

	override fun setRemoved() {
		super.setRemoved()
		unsubscribe()
	}

	override fun clearRemoved() {
		super.clearRemoved()

		if (level?.isClientSide == false) {
			subscribe()
		}
	}

	private var _subCache: ChunkSubscribers? = null
	private var playerListUpdated = false

	private fun unsubscribe() {
		val subCache = _subCache ?: return

		if (subCache.unsubscribe(this)) {
			val level = subCache.level.get()

			if (level != null) {
				playerMap.get(level)?.remove(subCache.chunkPos)
			}
		}

		playerListUpdated = false
	}

	private fun subscribe() {
		if (!syncher.hasChildren)
			return

		val level = level
		check(level is ServerLevel) { "Invalid realm" }
		unsubscribe()

		val subs = playerMap
			.computeIfAbsent(level) { Long2ObjectOpenHashMap() }
			.computeIfAbsent(ChunkPos(blockPos).toLong(), Long2ObjectFunction { ChunkSubscribers(level, ChunkPos(blockPos).toLong()) })

		subs.subscribe(this)
		_subCache = subs
		playerListUpdated = true
	}

	private class ChunkSubscribers(level: ServerLevel, val chunkPos: Long) {
		val level = WeakReference(level)
		val blockEntities = WeakHashSet<MatteryBlockEntity>(linked = true, initialCapacity = 0)
		val players = ObjectArraySet<ServerPlayer>(0)
		val veto = ObjectArraySet<ServerPlayer>(0)
		private val player2ResubCount = Reference2IntArrayMap<ServerPlayer>()
		private var isTicking = false

		private fun checkShouldTick() {
			val shouldTick = blockEntities.isNotEmpty() && players.isNotEmpty()

			if (!shouldTick && isTicking) {
				isTicking = false
				tickingMap[level.get()]?.remove(this)
			} else if (shouldTick && !isTicking) {
				isTicking = true
				tickingMap.computeIfAbsent(level.get()) { ArrayList() }.add(this)
			}
		}

		operator fun component1() = blockEntities
		operator fun component2() = players

		private fun recheckPlayer(player: ServerPlayer) {
			sometimeServer {
				if (player in players && !player.hasDisconnected()) {
					if (player.connection.chunkSender.isPending(chunkPos)) {
						recheckPlayer(player)
					} else {
						veto.remove(player)

						blockEntities.forEach {
							it.synchers[player] = it.syncher.Remote()
						}

						checkShouldTick()
					}
				} else if (player in players && player.hasDisconnected()) {
					unsubscribe(player)
				}
			}
		}

		fun subscribe(player: ServerPlayer) {
			if (players.add(player)) {
				veto.add(player)
				recheckPlayer(player)
				checkShouldTick()
			} else if (player !in veto) {
				player2ResubCount[player] = player2ResubCount.getInt(player) + 1
				LOGGER.debug("{} got subscribed to {} without prior unsubscribing, forcefully resubscribing (this happened {} times)", player, ChunkPos(chunkPos), player2ResubCount.getInt(player))
				unsubscribe(player, false)
				subscribe(player)
			}
		}

		fun unsubscribe(player: ServerPlayer, normal: Boolean = true): Boolean {
			if (players.remove(player)) {
				veto.remove(player)
				if (normal) player2ResubCount.removeInt(player)

				blockEntities.forEach {
					it.synchers.remove(player)
				}

				checkShouldTick()
				return true
			}

			return false
		}

		fun subscribe(blockEntity: MatteryBlockEntity) {
			if (blockEntities.add(blockEntity)) {
				players.forEach {
					if (it !in veto) {
						blockEntity.synchers[it] = blockEntity.syncher.Remote()
					}
				}

				checkShouldTick()
			}
		}

		fun unsubscribe(blockEntity: MatteryBlockEntity): Boolean {
			blockEntities.remove(blockEntity)
			checkShouldTick()
			return players.isEmpty() && blockEntities.isEmpty()
		}

		val isEmpty: Boolean get() {
			return players.isEmpty() && blockEntities.isEmpty()
		}
	}

	/**
	 * Why track player-tracked chunks?
	 *
	 * because minecraft itself doesn't track them, well,
	 * in "section partitioning" way.
	 *
	 * just look at [net.minecraft.server.level.PlayerMap]
	 *
	 * the [net.minecraft.server.level.PlayerMap.getPlayers] straight ignores chunk position
	 * and just fetches all players
	 *
	 * even if they did not ignore that argument, you still have to fetch player *list* though
	 * [net.minecraft.server.level.ChunkMap.getPlayers], which is not something we want
	 */
	companion object {
		const val REDSTONE_CONTROL_KEY = "redstoneControl"
		const val INVENTORY_KEY = "container"
		const val MATTER_STORAGE_KEY = "matterStorage"
		const val ENERGY_KEY = "energyStorage"
		const val BATTERY_KEY = "batteryInventory"
		const val LOOT_TABLE_KEY = RandomizableContainerBlockEntity.LOOT_TABLE_TAG
		const val LOOT_TABLE_SEED_KEY = RandomizableContainerBlockEntity.LOOT_TABLE_SEED_TAG

		private val playerMap = WeakHashMap<ServerLevel, Long2ObjectOpenHashMap<ChunkSubscribers>>()
		private val tickingMap = WeakHashMap<ServerLevel, ArrayList<ChunkSubscribers>>()

		private val vec2Dir = Int2ObjectOpenHashMap<Direction>()
		private val LOGGER = LogManager.getLogger()

		/**
		 * Returns stream of players watching (tracking) specified [chunkPos] in [level]
		 */
		fun watchingPlayers(chunkPos: Long, level: Level): Stream<ServerPlayer> {
			if (level !is ServerLevel)
				return Stream.empty()

			val subs = playerMap[level] ?: return Stream.empty()
			val chunk = subs[chunkPos] ?: return Stream.empty()
			return chunk.players.stream()
		}

		/**
		 * Returns stream of players watching (tracking) specified [chunkPos] in [level]
		 */
		fun watchingPlayers(chunkPos: ChunkPos, level: Level) = watchingPlayers(chunkPos.toLong(), level)

		/**
		 * Returns stream of players watching (tracking) specified [blockPos] in [level]
		 */
		fun watchingPlayers(blockPos: BlockPos, level: Level) = watchingPlayers(ChunkPos(blockPos), level)

		/**
		 * Returns stream of players watching (tracking) specified [blockPos] in [level]
		 */
		fun watchingPlayers(blockPos: Vec3, level: Level) = watchingPlayers(ChunkPos.asLong(SectionPos.blockToSectionCoord(blockPos.x), SectionPos.blockToSectionCoord(blockPos.z)), level)

		/**
		 * Returns stream of players watching (tracking) specified [chunk]
		 */
		fun watchingPlayers(chunk: LevelChunk) = watchingPlayers(chunk.pos, chunk.level!!)

		private fun vecKey(value: Vec3i): Int {
			if (value.x !in -1 .. 1) return -1
			if (value.y !in -1 .. 1) return -1
			if (value.z !in -1 .. 1) return -1
			val x = if (value.x < 0) 2 else value.x
			val y = if (value.y < 0) 2 else value.y
			val z = if (value.z < 0) 2 else value.z
			return x or (y shl 4) or (z shl 8)
		}

		init {
			for (dir in Direction.entries) {
				vec2Dir[vecKey(dir.normal)] = dir
			}
		}

		fun onLevelUnload(event: LevelEvent.Unload) {
			val level = event.level as? ServerLevel ?: return
			playerMap.remove(level)
			tickingMap.remove(level)
		}

		fun onServerStopping(event: ServerStoppingEvent) {
			playerMap.clear()
			tickingMap.clear()
		}

		fun postLevelTick(event: LevelTickEvent.Post) {
			val level = event.level as? ServerLevel ?: return

			tickingMap[level]?.forEach {
				it.blockEntities.forEach { be ->
					be.syncher.observe()

					be.synchers.entries.forEach { e ->
						val (player, data) = e
						val payload = data.write(event.level.registryAccess())

						if (payload != null) {
							PacketDistributor.sendToPlayer(player, BlockEntitySyncPacket(be.blockPos, payload))
						}
					}
				}
			}
		}

		fun onWatch(event: ChunkWatchEvent.Watch) {
			playerMap
				.computeIfAbsent(event.level) { Long2ObjectOpenHashMap() }
				.computeIfAbsent(event.pos.toLong(), Long2ObjectFunction { ChunkSubscribers(event.level, it) })
				.subscribe(event.player)
		}

		fun onForget(event: ChunkWatchEvent.UnWatch) {
			val levelMap = playerMap[event.level] ?: return
			val subs = levelMap.get(event.pos.toLong()) ?: return

			if (subs.unsubscribe(event.player) && subs.isEmpty) {
				levelMap.remove(event.pos.toLong())
			}
		}

		fun playerDisconnected(event: PlayerEvent.PlayerLoggedOutEvent) {
			for (tree in playerMap.values) {
				tree.values.removeIf {
					it.unsubscribe(event.entity as ServerPlayer)
					it.isEmpty
				}
			}
		}
	}
}
