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

package net.shadowfacts.phycon.block

import net.minecraft.block.BlockState
import net.minecraft.block.entity.BlockEntity
import net.minecraft.block.entity.BlockEntityType
import net.minecraft.nbt.NbtCompound
import net.minecraft.network.packet.s2c.play.BlockEntityUpdateS2CPacket
import net.minecraft.util.math.BlockPos
import net.shadowfacts.phycon.PhysicalConnectivity
import net.shadowfacts.phycon.api.Interface
import net.shadowfacts.phycon.api.NetworkComponentBlock
import net.shadowfacts.phycon.api.PacketSink
import net.shadowfacts.phycon.api.PacketSource
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.frame.ARPQueryFrame
import net.shadowfacts.phycon.frame.ARPResponseFrame
import net.shadowfacts.phycon.frame.BasePacketFrame
import net.shadowfacts.phycon.packet.*
import net.shadowfacts.phycon.util.NetworkUtil
import java.lang.ref.WeakReference
import java.util.*

/**
 * @author shadowfacts
 */
abstract class DeviceBlockEntity(type: BlockEntityType<*>, pos: BlockPos, state: BlockState): BlockEntity(type, pos, state),
	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>()
	private var cachedDestination: WeakReference<Interface>? = null

	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)
			}
			is PingPacket -> {
				sendPacket(PongPacket(ipAddress, 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 cachedDestination = this.cachedDestination?.get()
		if (cachedDestination != null) {
			return cachedDestination
		}

		val sides = (cachedState.block as NetworkComponentBlock).getNetworkConnectedSides(cachedState, world!!, pos)
		return when (sides.size) {
			0 -> null
			1 -> {
				NetworkUtil.findConnectedInterface(world!!, pos, sides.first())?.also {
					this.cachedDestination = WeakReference(it)
				}
			}
			else -> throw RuntimeException("DeviceBlockEntity.findDestination must be overridden by devices which have more than 1 network connected side")
		}
	}

	override fun cableDisconnected() {
		cachedDestination = null
	}

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

	protected open fun toCommonTag(tag: NbtCompound) {
		tag.putInt("IPAddress", ipAddress.address)
		tag.putLong("MACAddress", macAddress.address)
	}

	protected open fun fromCommonTag(tag: NbtCompound) {
		ipAddress = IPAddress(tag.getInt("IPAddress"))
		macAddress = MACAddress(tag.getLong("MACAddress"))
	}

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

	override fun readNbt(tag: NbtCompound) {
		super.readNbt(tag)
		fromCommonTag(tag)
		if (tag.getBoolean("_SyncPacket")) {
			fromClientTag(tag)
		}
	}

	override fun toUpdatePacket(): BlockEntityUpdateS2CPacket {
		return BlockEntityUpdateS2CPacket.create(this)
	}

	override fun toInitialChunkDataNbt(): NbtCompound {
		val tag = NbtCompound()
		tag.putBoolean("_SyncPacket", true)
		return toClientTag(tag)
	}

	open fun toClientTag(tag: NbtCompound): NbtCompound {
		toCommonTag(tag)
		return tag
	}

	open fun fromClientTag(tag: NbtCompound) {
	}

	fun markUpdate() {
		markDirty()
		world!!.updateListeners(pos, cachedState, cachedState, 3)
	}

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

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

}