PhysicalConnectivity/src/main/kotlin/net/shadowfacts/phycon/block/netswitch/SwitchBlockEntity.kt

package net.shadowfacts.phycon.block.netswitch

import net.minecraft.block.BlockState
import net.minecraft.block.entity.BlockEntity
import net.minecraft.item.ItemStack
import net.minecraft.nbt.NbtCompound
import net.minecraft.nbt.NbtList
import net.minecraft.network.Packet
import net.minecraft.network.listener.ClientPlayPacketListener
import net.minecraft.network.packet.s2c.play.BlockEntityUpdateS2CPacket
import net.minecraft.util.math.BlockPos
import net.minecraft.util.math.Direction
import net.shadowfacts.phycon.PhysicalConnectivity
import net.shadowfacts.phycon.api.Interface
import net.shadowfacts.phycon.api.frame.EthernetFrame
import net.shadowfacts.phycon.api.frame.PacketFrame
import net.shadowfacts.phycon.api.util.IPAddress
import net.shadowfacts.phycon.api.util.MACAddress
import net.shadowfacts.phycon.frame.BasePacketFrame
import net.shadowfacts.phycon.init.PhyBlockEntities
import net.shadowfacts.phycon.packet.ItemStackPacket
import net.shadowfacts.phycon.util.NetworkUtil
import java.lang.ref.WeakReference
import java.util.Deque
import java.util.LinkedList

/**
 * @author shadowfacts
 */
class SwitchBlockEntity(pos: BlockPos, state: BlockState): BlockEntity(PhyBlockEntities.SWITCH, pos, state) {

	companion object {
		var SWITCHING_CAPACITY = 256 // 256 packets/tick
	}

	val interfaces = Direction.values().map { SwitchInterface(it, WeakReference(this), MACAddress.random()) }

	private val macTable = mutableMapOf<MACAddress, Direction>()
	private var packetsHandledThisTick = 0
	private var delayedPackets: Deque<Pair<PacketFrame, SwitchInterface>> = LinkedList()

	fun interfaceForSide(side: Direction): SwitchInterface {
		return interfaces.find { it.side == side }!!
	}

	private fun handle(frame: EthernetFrame, fromItf: SwitchInterface) {
		macTable[frame.source] = fromItf.side

		if (frame is PacketFrame) {
			if (packetsHandledThisTick > SWITCHING_CAPACITY) {
				PhysicalConnectivity.NETWORK_LOGGER.debug("{} reached capacity, delaying forwarding {}", this, frame)
				delayedPackets.addLast(frame to fromItf)
				return
			} else {
				packetsHandledThisTick++
			}
		}

		resend(frame, fromItf)
	}

	private fun resend(frame: EthernetFrame, fromItf: SwitchInterface) {
		if (frame.destination.type != MACAddress.Type.BROADCAST && macTable.containsKey(frame.destination)) {
			val dir = macTable[frame.destination]!!
			PhysicalConnectivity.NETWORK_LOGGER.debug("{} ({}, {}) forwarding {} to side {}", this, fromItf.side, fromItf.macAddress, frame, dir)
			interfaceForSide(dir).send(frame)
		} else {
			PhysicalConnectivity.NETWORK_LOGGER.debug("{} ({}, {}) flooding {}", this, fromItf.side, fromItf.macAddress, frame)
			flood(frame, fromItf)
		}
	}

	private fun flood(frame: EthernetFrame, source: Interface) {
		for (itf in interfaces) {
			if (source == itf) continue
			itf.send(frame)
		}
	}

	private fun findDestination(fromItf: Interface): Interface? {
		val side = (fromItf as SwitchInterface).side
		return NetworkUtil.findConnectedInterface(world!!, pos, side)
	}

	fun tick() {
		packetsHandledThisTick = 0

		while (delayedPackets.isNotEmpty() && packetsHandledThisTick <= SWITCHING_CAPACITY) {
			val (frame, fromItf) = delayedPackets.pop()
			resend(frame, fromItf)
		}
	}

	override fun writeNbt(tag: NbtCompound) {
		super.writeNbt(tag)

		tag.putLongArray("InterfaceAddresses", interfaces.map { it.macAddress.address })
		val list = NbtList()
		for ((frame, fromItf) in delayedPackets) {
			val packet = frame.packet
			if (packet !is ItemStackPacket) continue
			val compound = NbtCompound()
			compound.putInt("FromItfSide", fromItf.side.ordinal)
			compound.putInt("SourceIP", packet.source.address)
			compound.putInt("DestinationIP", packet.destination.address)
			compound.putLong("SourceMAC", frame.source.address)
			compound.putLong("DestinationMAC", frame.destination.address)
			compound.put("Stack", packet.stack.writeNbt(NbtCompound()))
			list.add(compound)
		}
		tag.put("DelayedStackPackets", list)
	}

	override fun readNbt(tag: NbtCompound) {
		super.readNbt(tag)

		tag.getLongArray("InterfaceAddresses")?.forEachIndexed { i, l ->
			interfaces[i].macAddress = MACAddress(l)
		}
		tag.getList("DelayedStackPackets", 10).forEach { it ->
			val compound = it as NbtCompound
			val fromItfSide = Direction.values()[compound.getInt("FromItfSide")]
			val fromItf = interfaces.find { it.side == fromItfSide }!!
			val sourceIP = IPAddress(compound.getInt("SourceIP"))
			val destinationIP = IPAddress(compound.getInt("DestinationIP"))
			val sourceMAC = MACAddress(compound.getLong("SourceMAC"))
			val destinationMAC = MACAddress(compound.getLong("DestinationMAC"))
			val stack = ItemStack.fromNbt(compound.getCompound("Stack"))
			if (!stack.isEmpty) {
				val packet = ItemStackPacket(stack, sourceIP, destinationIP)
				val frame = BasePacketFrame(packet, sourceMAC, destinationMAC)
				delayedPackets.addLast(frame to fromItf)
			}
		}
	}

	override fun toUpdatePacket(): Packet<ClientPlayPacketListener>? {
		return BlockEntityUpdateS2CPacket.create(this)
	}

	override fun toInitialChunkDataNbt(): NbtCompound {
		val tag = NbtCompound()
		tag.putLongArray("InterfaceAddresses", interfaces.map { it.macAddress.address })
		return tag
	}

	class SwitchInterface(
		val side: Direction,
		val switch: WeakReference<SwitchBlockEntity>,
		@JvmField var macAddress: MACAddress,
	): Interface {
		override fun getMACAddress() = macAddress

		override fun receive(frame: EthernetFrame) {
			switch.get()?.handle(frame, this)
		}

		override fun send(frame: EthernetFrame) {
			switch.get()?.findDestination(this)?.receive(frame)
		}
	}

}