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

package net.shadowfacts.phycon.block.netswitch

import net.fabricmc.fabric.api.block.entity.BlockEntityClientSerializable
import net.minecraft.block.BlockState
import net.minecraft.block.entity.BlockEntity
import net.minecraft.item.ItemStack
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.ListTag
import net.minecraft.util.Tickable
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: BlockEntity(PhyBlockEntities.SWITCH),
	BlockEntityClientSerializable,
	Tickable {

	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)
	}

	override fun tick() {
		packetsHandledThisTick = 0

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

	override fun toTag(tag: CompoundTag): CompoundTag {
		tag.putLongArray("InterfaceAddresses", interfaces.map { it.macAddress.address })
		val list = ListTag()
		for ((frame, fromItf) in delayedPackets) {
			val packet = frame.packet
			if (packet !is ItemStackPacket) continue
			val compound = CompoundTag()
			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.toTag(CompoundTag()))
			list.add(compound)
		}
		tag.put("DelayedStackPackets", list)
		return super.toTag(tag)
	}

	override fun fromTag(state: BlockState, tag: CompoundTag) {
		super.fromTag(state, tag)
		tag.getLongArray("InterfaceAddresses")?.forEachIndexed { i, l ->
			interfaces[i].macAddress = MACAddress(l)
		}
		tag.getList("DelayedStackPackets", 10).forEach { it ->
			val compound = it as CompoundTag
			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.fromTag(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 toClientTag(tag: CompoundTag): CompoundTag {
		tag.putLongArray("InterfaceAddresses", interfaces.map { it.macAddress.address })
		return tag
	}

	override fun fromClientTag(tag: CompoundTag) {
		tag.getLongArray("InterfaceAddresses")?.forEachIndexed { i, l ->
			interfaces[i].macAddress = MACAddress(l)
		}
	}

	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)
		}
	}

}