PhysicalConnectivity/src/main/kotlin/net/shadowfacts/phycon/network/block/terminal/TerminalBlockEntity.kt

package net.shadowfacts.phycon.network.block.terminal

import alexiil.mc.lib.attributes.item.GroupedItemInv
import alexiil.mc.lib.attributes.item.ItemStackCollections
import alexiil.mc.lib.attributes.item.ItemStackUtil
import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap
import net.fabricmc.fabric.api.block.entity.BlockEntityClientSerializable
import net.fabricmc.fabric.api.screenhandler.v1.ExtendedScreenHandlerFactory
import net.minecraft.block.BlockState
import net.minecraft.entity.player.PlayerEntity
import net.minecraft.entity.player.PlayerInventory
import net.minecraft.inventory.Inventory
import net.minecraft.inventory.InventoryChangedListener
import net.minecraft.item.ItemStack
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.ListTag
import net.minecraft.network.PacketByteBuf
import net.minecraft.screen.ScreenHandler
import net.minecraft.server.network.ServerPlayerEntity
import net.minecraft.text.LiteralText
import net.minecraft.util.Tickable
import net.minecraft.util.math.Direction
import net.shadowfacts.phycon.api.Interface
import net.shadowfacts.phycon.api.packet.Packet
import net.shadowfacts.phycon.api.util.IPAddress
import net.shadowfacts.phycon.init.PhyBlockEntities
import net.shadowfacts.phycon.network.DeviceBlockEntity
import net.shadowfacts.phycon.network.NetworkUtil
import net.shadowfacts.phycon.network.block.netinterface.InterfaceBlockEntity
import net.shadowfacts.phycon.network.component.ItemStackPacketHandler
import net.shadowfacts.phycon.network.component.handleItemStack
import net.shadowfacts.phycon.network.packet.*
import java.util.*
import kotlin.math.min

/**
 * @author shadowfacts
 */
class TerminalBlockEntity: DeviceBlockEntity(PhyBlockEntities.TERMINAL), InventoryChangedListener, BlockEntityClientSerializable, Tickable, ItemStackPacketHandler {

	companion object {
		val LOCATE_REQUEST_TIMEOUT = 40 // ticks
        val INSERTION_TIMEOUT = 40
	}

	private val inventoryCache = mutableMapOf<IPAddress, GroupedItemInv>()
	val internalBuffer = TerminalBufferInventory(18)

	private val locateRequestQueue = LinkedList<StackLocateRequest>()
	private val pendingRequests = LinkedList<StackLocateRequest>()
    private val pendingInsertions = Int2ObjectArrayMap<PendingStackInsertion>()

	private var observers = 0
	val cachedNetItems = ItemStackCollections.intMap()
	var cachedSortedNetItems = listOf<ItemStack>()

	init {
		internalBuffer.addListener(this)
	}

	override fun findDestination(): Interface? {
		for (dir in Direction.values()) {
			val itf = NetworkUtil.findConnectedInterface(world!!, pos, dir)
			if (itf != null) {
				return itf
			}
		}
		return null
	}

	override fun handle(packet: Packet) {
		when (packet) {
			is ReadInventoryPacket -> handleReadInventory(packet)
			is DeviceRemovedPacket -> handleDeviceRemoved(packet)
			is StackLocationPacket -> handleStackLocation(packet)
			is ItemStackPacket -> handleItemStack(packet)
			is CapacityPacket -> handleCapacity(packet)
		}
	}

	private fun handleReadInventory(packet: ReadInventoryPacket) {
		inventoryCache[packet.source] = packet.inventory
		updateNetItems()
		sync()
	}

	private fun handleDeviceRemoved(packet: DeviceRemovedPacket) {
		inventoryCache.remove(packet.source)
		updateNetItems()
		sync()
	}

	private fun handleStackLocation(packet: StackLocationPacket) {
		val request = pendingRequests.firstOrNull {
			ItemStackUtil.areEqualIgnoreAmounts(it.stack, packet.stack)
		}
		if (request != null) {
			request.results.add(packet.amount to packet.sourceInterface)
			if (request.totalResultAmount >= request.amount || counter - request.timestamp >= LOCATE_REQUEST_TIMEOUT || request.results.size >= inventoryCache.size) {
				stackLocateRequestCompleted(request)
			}
		}
	}

	override fun doHandleItemStack(packet: ItemStackPacket): ItemStack {
		val remaining = internalBuffer.insertFromNetwork(packet.stack)

		// this happens outside the normal update loop because by receiving the item stack packet
		// we "know" how much the count in the source inventory has changed
		updateNetItems()
		sync()

		return remaining
	}

	private fun handleCapacity(packet: CapacityPacket) {
		val insertion = pendingInsertions.values.firstOrNull {
			ItemStackUtil.areEqualIgnoreAmounts(packet.stack, it.stack)
		}
		if (insertion != null) {
			insertion.results.add(packet.capacity to packet.receivingInterface)
			if (insertion.isFinishable(counter)) {
				finishInsertion(insertion)
			}
		}
	}

	private fun updateNetItems() {
		cachedNetItems.clear()
		for (inventory in inventoryCache.values) {
			for (stack in inventory.storedStacks) {
				val amount = inventory.getAmount(stack)
				cachedNetItems.mergeInt(stack, amount) { a, b -> a + b }
			}
		}
		// todo: is the map necessary or is just the sorted list enough?
		cachedSortedNetItems = cachedNetItems.object2IntEntrySet().sortedByDescending { it.intValue }.map {
			val stack = it.key.copy()
			stack.count = it.intValue
			stack
		}
	}

	private fun beginInsertions() {
		if (world!!.isClient) return

		for (slot in 0 until internalBuffer.size()) {
			if (internalBuffer.getMode(slot) != TerminalBufferInventory.Mode.TO_NETWORK) continue
			if (slot in pendingInsertions) continue
			val stack = internalBuffer.getStack(slot)
			pendingInsertions[slot] = PendingStackInsertion(slot, stack, counter)
			sendPacket(CheckCapacityPacket(stack, ipAddress, IPAddress.BROADCAST))
		}
	}

	private fun finishPendingInsertions() {
		if (world!!.isClient) return

		for (insertion in pendingInsertions.values) {
			if (!insertion.isFinishable(counter)) continue
			finishInsertion(insertion)
		}
	}

	private fun sendEnqueuedLocateRequests() {
		if (world!!.isClient) return

		for (request in locateRequestQueue) {
			pendingRequests.add(request)
			sendPacket(LocateStackPacket(request.stack, ipAddress))
		}
		locateRequestQueue.clear()
	}

	private fun finishPendingRequests() {
		if (world!!.isClient) return

		for (request in pendingRequests) {
			if (request.isFinishable(counter)) {
				stackLocateRequestCompleted(request)
			}
		}
	}

	fun addObserver() {
		observers++
	}

	fun removeObserver() {
		observers--
	}

	override fun tick() {
		super.tick()

		if (counter % 20 == 0L) {
			if (!world!!.isClient) {
				sendEnqueuedLocateRequests()
				finishPendingRequests()
				beginInsertions()
				finishPendingInsertions()
			}

			if (observers > 0) {
				if (world!!.isClient) {
					println(cachedNetItems)
				} else {
					updateNetItems()
					sync()
				}
			}
		}
	}

	fun onActivate(player: PlayerEntity) {
		if (!world!!.isClient) {
			updateNetItems()
			sync()

			inventoryCache.clear()
			sendPacket(RequestInventoryPacket(ipAddress))
			val factory = object: ExtendedScreenHandlerFactory {
				override fun createMenu(syncId: Int, playerInv: PlayerInventory, player: PlayerEntity): ScreenHandler? {
					return TerminalScreenHandler(syncId, playerInv, this@TerminalBlockEntity)
				}

				override fun getDisplayName() = LiteralText("Terminal")

				override fun writeScreenOpeningData(player: ServerPlayerEntity, buf: PacketByteBuf) {
					buf.writeBlockPos(this@TerminalBlockEntity.pos)
				}
			}
			player.openHandledScreen(factory)
		}
		addObserver()
	}

	fun requestItem(stack: ItemStack, amount: Int = stack.count) {
		locateRequestQueue.add(StackLocateRequest(stack, amount, counter))
	}

	private fun stackLocateRequestCompleted(request: StackLocateRequest) {
		pendingRequests.remove(request)

		// todo: also sort results by interface priority
		val sortedResults = request.results.sortedByDescending { it.first }.toMutableList()
		var amountRequested = 0
		while (amountRequested < request.amount && sortedResults.isNotEmpty()) {
			val (sourceAmount, sourceInterface) = sortedResults.removeAt(0)
			val amountToRequest = min(sourceAmount, request.amount - amountRequested)
			amountRequested += amountToRequest
			sendPacket(ExtractStackPacket(request.stack, amountToRequest, ipAddress, sourceInterface.ipAddress))
		}
	}

	private fun finishInsertion(insertion: PendingStackInsertion) {
		pendingInsertions.remove(insertion.bufferSlot)

		// todo: also sort results by interface priority
		val sortedResults = insertion.results.sortedBy { it.first }.toMutableList()
		val remaining = insertion.stack
		while (!remaining.isEmpty && sortedResults.isNotEmpty()) {
			val (capacity, receivingInterface) = sortedResults.removeAt(0)
			if (capacity <= 0) continue
			sendPacket(ItemStackPacket(remaining.copy(), ipAddress, receivingInterface.ipAddress))
			// todo: the interface should confirm how much was actually inserted, in case of race condition
			remaining.count -= capacity
		}
        internalBuffer.setStack(insertion.bufferSlot, remaining)

		// as with extracting, we "know" the new amounts and so can update instantly without actually sending out packets
		updateNetItems()
		sync()
	}

	override fun onInventoryChanged(inv: Inventory) {
		if (inv == internalBuffer && world != null && !world!!.isClient) {
			markDirty()
			sync()
		}
	}

	override fun toTag(tag: CompoundTag): CompoundTag {
		tag.put("InternalBuffer", internalBuffer.toTag())
		return super.toTag(tag)
	}

	override fun fromTag(state: BlockState, tag: CompoundTag) {
		super.fromTag(state, tag)
		internalBuffer.fromTag(tag.getCompound("InternalBuffer"))
	}

	override fun toClientTag(tag: CompoundTag): CompoundTag {
		tag.put("InternalBuffer", internalBuffer.toTag())
		val list = ListTag()
		tag.put("CachedNetItems", list)
		for ((stack, amount) in cachedNetItems) {
			val entryTag = stack.toTag(CompoundTag())
			entryTag.putInt("NetAmount", amount)
			list.add(entryTag)
		}
		return tag
	}

	override fun fromClientTag(tag: CompoundTag) {
		internalBuffer.fromTag(tag.getCompound("InternalBuffer"))
		val list = tag.getList("CachedNetItems", 10)
		cachedNetItems.clear()
		for (entryTag in list) {
			val stack = ItemStack.fromTag(entryTag as CompoundTag)
			val netAmount = entryTag.getInt("NetAmount")
			cachedNetItems[stack] = netAmount
		}
		cachedSortedNetItems = cachedNetItems.object2IntEntrySet().sortedByDescending { it.intValue }.map {
			val stack = it.key.copy()
			stack.count = it.intValue
			stack
		}
	}

}

data class StackLocateRequest(
		val stack: ItemStack,
		val amount: Int,
		val timestamp: Long,
		var results: MutableSet<Pair<Int, InterfaceBlockEntity>> = mutableSetOf()
) {
	val totalResultAmount: Int
		get() = results.fold(0) { acc, (amount, _) -> acc + amount }

	fun isFinishable(currentTimestamp: Long): Boolean {
		return totalResultAmount >= amount || currentTimestamp - timestamp >= TerminalBlockEntity.LOCATE_REQUEST_TIMEOUT
	}
}

data class PendingStackInsertion(
	val bufferSlot: Int,
	val stack: ItemStack,
	val timestamp: Long,
	val results: MutableSet<Pair<Int, InterfaceBlockEntity>> = mutableSetOf(),
) {
	val totalCapacity: Int
		get() = results.fold(0) { acc, (amount, _) -> acc + amount }

	fun isFinishable(currentTimestamp: Long): Boolean {
		return totalCapacity >= stack.count || currentTimestamp - timestamp >= TerminalBlockEntity.INSERTION_TIMEOUT
	}
}