PhysicalConnectivity/src/main/kotlin/net/shadowfacts/phycon/network/DeviceBlockEntity.kt

package net.shadowfacts.phycon.network

import net.fabricmc.fabric.api.block.entity.BlockEntityClientSerializable
import net.minecraft.block.BlockState
import net.minecraft.block.entity.BlockEntity
import net.minecraft.block.entity.BlockEntityType
import net.minecraft.nbt.CompoundTag
import net.minecraft.util.Tickable
import net.shadowfacts.phycon.PhysicalConnectivity
import net.shadowfacts.phycon.api.PacketSink
import net.shadowfacts.phycon.api.PacketSource
import net.shadowfacts.phycon.api.Interface
import net.shadowfacts.phycon.api.NetworkComponentBlock
import net.shadowfacts.phycon.api.frame.EthernetFrame
import net.shadowfacts.phycon.api.frame.PacketFrame
import net.shadowfacts.phycon.api.packet.Packet
import net.shadowfacts.phycon.api.util.IPAddress
import net.shadowfacts.phycon.api.util.MACAddress
import net.shadowfacts.phycon.network.frame.ARPQueryFrame
import net.shadowfacts.phycon.network.frame.ARPResponseFrame
import net.shadowfacts.phycon.network.frame.BasePacketFrame
import net.shadowfacts.phycon.network.packet.*
import java.util.*

/**
 * @author shadowfacts
 */
abstract class DeviceBlockEntity(type: BlockEntityType<*>): BlockEntity(type),
	BlockEntityClientSerializable,
	Tickable,
	PacketSink,
	PacketSource,
	Interface {

	companion object {
		private const val ARP_RETRY_TIMEOUT = 200
	}

	var macAddress: MACAddress = MACAddress.random()
		protected set

	var ipAddress: IPAddress = IPAddress.random()
		protected set

	private val arpTable = mutableMapOf<IPAddress, MACAddress>()
	private val packetQueue = LinkedList<PendingPacket>()

	var counter: Long = 0

	override fun getIPAddress() = ipAddress
	override fun getMACAddress() = macAddress

	abstract override fun handle(packet: Packet)

	private fun doHandlePacket(packet: Packet) {
		PhysicalConnectivity.NETWORK_LOGGER.debug("{} ({}) received packet: {}", this, ipAddress, packet)
		when (packet) {
			is DeviceRemovedPacket -> {
				arpTable.remove(packet.source)
			}
		}
		handle(packet)
	}

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

	override fun receive(frame: EthernetFrame) {
		PhysicalConnectivity.NETWORK_LOGGER.debug("{} ({}, {}) received frame from {}: {}", this, ipAddress, macAddress, frame.source, frame)
		when (frame) {
			is ARPQueryFrame -> handleARPQuery(frame)
			is ARPResponseFrame -> handleARPResponse(frame)
			is PacketFrame -> {
				if (frame.packet.destination.isBroadcast || frame.packet.destination == ipAddress) {
					doHandlePacket(frame.packet)
				}
			}
		}
	}

	private fun handleARPQuery(frame: ARPQueryFrame) {
		PhysicalConnectivity.NETWORK_LOGGER.debug("{}, ({}), received ARP query for {}", this, ipAddress, frame.queryIP)
		arpTable[frame.sourceIP] = frame.source
		if (frame.queryIP == ipAddress) {
			PhysicalConnectivity.NETWORK_LOGGER.debug("{} ({}) sending ARP response to {} with {}", this, ipAddress, frame.sourceIP, macAddress)
			send(ARPResponseFrame(ipAddress, macAddress, frame.source))
		}
	}

	private fun handleARPResponse(frame: ARPResponseFrame) {
		arpTable[frame.query] = frame.source
		PhysicalConnectivity.NETWORK_LOGGER.debug("{}, ({}) received ARP response for {} with {}", this, ipAddress, frame.query, frame.source)

		val toRemove = packetQueue.filter { (packet, _) ->
			if (packet.destination == frame.query) {
				send(BasePacketFrame(packet, macAddress, frame.source))
				true
			} else {
				false
			}
		}
		packetQueue.removeAll(toRemove)
	}

	override fun sendPacket(packet: Packet) {
		if (packet.destination.isBroadcast) {
			send(BasePacketFrame(packet, macAddress, MACAddress.BROADCAST))
		} else if (arpTable.containsKey(packet.destination)) {
			send(BasePacketFrame(packet, macAddress, arpTable[packet.destination]!!))
		} else {
			packetQueue.add(PendingPacket(packet, counter))

			PhysicalConnectivity.NETWORK_LOGGER.debug("{} ({}) sending ARP query for {}", this, ipAddress, packet.destination)
			send(ARPQueryFrame(packet.destination, ipAddress, macAddress))
		}
	}

	open fun findDestination(): Interface? {
		val sides = (cachedState.block as NetworkComponentBlock).getNetworkConnectedSides(cachedState, world!!, pos)
		if (sides.size != 1) {
			throw RuntimeException("DeviceBlockEntity.findDestination must be overridden by devices which have more than 1 network connected side")
		}
		return NetworkUtil.findConnectedInterface(world!!, pos, sides.first())
	}

	override fun tick() {
		counter++

		if (!world!!.isClient) {
			val toRemove = packetQueue.filter { entry ->
				val (packet, timestamp) = entry
				if (arpTable.containsKey(packet.destination)) {
					send(BasePacketFrame(packet, macAddress, arpTable[packet.destination]!!))
					true
				} else if (counter - timestamp >= ARP_RETRY_TIMEOUT) {
					send(ARPQueryFrame(packet.destination, ipAddress, macAddress))
					entry.timestamp = counter
					// todo: should there be a retry counter?
					true
				} else {
					false
				}
			}
			packetQueue.removeAll(toRemove)
		}
	}

	override fun toTag(tag: CompoundTag): CompoundTag {
		tag.putInt("IPAddress", ipAddress.address)
		tag.putLong("MACAddress", macAddress.address)
		return super.toTag(tag)
	}

	override fun fromTag(state: BlockState, tag: CompoundTag) {
		super.fromTag(state, tag)
		ipAddress = IPAddress(tag.getInt("IPAddress"))
		macAddress = MACAddress(tag.getLong("MACAddress"))
	}

	override fun toClientTag(tag: CompoundTag): CompoundTag {
		tag.putInt("IPAddress", ipAddress.address)
		tag.putLong("MACAddress", macAddress.address)
		return tag
	}

	override fun fromClientTag(tag: CompoundTag) {
		ipAddress = IPAddress(tag.getInt("IPAddress"))
		macAddress = MACAddress(tag.getLong("MACAddress"))
	}

	fun onBreak() {
		if (!world!!.isClient) {
			sendPacket(DeviceRemovedPacket(this))
		}
	}

	data class PendingPacket(
		val packet: Packet,
		var timestamp: Long,
	)

}