package ru.dbotthepony.kstarbound.network

import io.netty.buffer.ByteBuf
import io.netty.buffer.ByteBufOutputStream
import io.netty.channel.ChannelDuplexHandler
import io.netty.channel.ChannelHandlerContext
import io.netty.channel.ChannelPromise
import it.unimi.dsi.fastutil.bytes.ByteArrayList
import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap
import it.unimi.dsi.fastutil.io.FastByteArrayInputStream
import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.kommons.io.readSignedVarInt
import ru.dbotthepony.kommons.io.writeSignedVarInt
import ru.dbotthepony.kstarbound.client.network.packets.ForgetChunkPacket
import ru.dbotthepony.kstarbound.client.network.packets.ChunkCellsPacket
import ru.dbotthepony.kstarbound.client.network.packets.ForgetEntityPacket
import ru.dbotthepony.kstarbound.client.network.packets.JoinWorldPacket
import ru.dbotthepony.kstarbound.client.network.packets.SpawnWorldObjectPacket
import ru.dbotthepony.kstarbound.network.packets.ClientContextUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.EntityCreatePacket
import ru.dbotthepony.kstarbound.network.packets.EntityDestroyPacket
import ru.dbotthepony.kstarbound.network.packets.EntityUpdateSetPacket
import ru.dbotthepony.kstarbound.network.packets.PingPacket
import ru.dbotthepony.kstarbound.network.packets.PongPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.ClientConnectPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.ConnectSuccessPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.HandshakeChallengePacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.HandshakeResponsePacket
import ru.dbotthepony.kstarbound.network.packets.ProtocolRequestPacket
import ru.dbotthepony.kstarbound.network.packets.ProtocolResponsePacket
import ru.dbotthepony.kstarbound.network.packets.StepUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.CelestialResponsePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.CentralStructureUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.ChatReceivePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.ConnectFailurePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.EntityInteractResultPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.EnvironmentUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.FindUniqueEntityResponsePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.GiveItemPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.LegacyTileArrayUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.LegacyTileUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.PlayerWarpResultPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.ServerDisconnectPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.ServerInfoPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SetPlayerStartPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SystemObjectCreatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SystemObjectDestroyPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SystemShipCreatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SystemShipDestroyPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SystemWorldStartPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.SystemWorldUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.TileDamageUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.TileModificationFailurePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.UniverseTimeUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.UpdateWorldPropertiesPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.WorldStartPacket
import ru.dbotthepony.kstarbound.network.packets.clientbound.WorldStopPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.CelestialRequestPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.ChatSendPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.ClientDisconnectRequestPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.ConnectWirePacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.DamageTileGroupPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.EntityInteractPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.FindUniqueEntityPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.FlyShipPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.ModifyTileListPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.PlayerWarpPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.RequestDropPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.SpawnEntityPacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.WorldClientStateUpdatePacket
import ru.dbotthepony.kstarbound.network.packets.serverbound.WorldStartAcknowledgePacket
import java.io.BufferedInputStream
import java.io.DataInputStream
import java.io.DataOutputStream
import java.io.EOFException
import java.io.FilterInputStream
import java.io.InputStream
import java.util.zip.Deflater
import java.util.zip.InflaterInputStream
import kotlin.math.absoluteValue
import kotlin.reflect.KClass

class PacketRegistry(val isLegacy: Boolean) {
	private val packets = ArrayList<Type<*>?>()
	private val missingNames = Int2ObjectArrayMap<String>()
	private val clazz2Type = Reference2ObjectOpenHashMap<KClass<*>, Type<*>>()

	private data class Type<T : IPacket>(val id: Int, val type: KClass<T>, val factory: IPacketReaderDetailed<T>, val direction: PacketDirection)

	val size: Int
		get() = packets.size

	private fun <T : IPacket> add(type: KClass<T>, reader: IPacketReaderDetailed<T>, direction: PacketDirection = PacketDirection.get(type)): PacketRegistry {
		if (packets.size >= 255)
			throw IndexOutOfBoundsException("Unable to add any more packet types! 255 is the max")

		if (type in clazz2Type)
			throw IllegalArgumentException("Already has packet handler for type $reader (${type})")

		val ptype = Type(packets.size, type, reader, direction)
		packets.add(ptype)
		clazz2Type[type] = ptype
		return this
	}

	private inline fun <reified T : IPacket> add(reader: IPacketReader<T>, direction: PacketDirection = PacketDirection.get(T::class)): PacketRegistry {
		return add(T::class, reader, direction)
	}

	private inline fun <reified T : IPacket> add(reader: IPacketReaderDetailed<T>, direction: PacketDirection = PacketDirection.get(T::class)): PacketRegistry {
		return add(T::class, reader, direction)
	}

	private inline fun <reified T : IPacket> add(value: T, direction: PacketDirection = PacketDirection.get(T::class)): PacketRegistry {
		return add(T::class, { _, _, _ -> value }, direction)
	}

	private fun skip(amount: Int = 1) {
		for (i in 0 until amount) {
			packets.add(null)
		}
	}

	private fun skip(name: String) {
		missingNames[packets.size] = name
		packets.add(null)
	}

	// avoid zip bomb
	private class LimitingInputStream(inputStream: InputStream) : FilterInputStream(inputStream) {
		private var read = 0L

		override fun read(): Int {
			if (read >= MAX_PACKET_SIZE)
				return -1

			read++
			return super.read()
		}

		override fun read(b: ByteArray, off: Int, len: Int): Int {
			if (read >= MAX_PACKET_SIZE)
				return -1

			val actual = super.read(b, off, len.coerceAtMost((MAX_PACKET_SIZE - read).coerceAtMost(Int.MAX_VALUE.toLong()).toInt()))

			if (actual > 0)
				read += actual

			return actual
		}

		override fun available(): Int {
			if (read >= MAX_PACKET_SIZE)
				return 0
			else
				return super.available()
		}
	}

	inner class Serializer(val side: ConnectionSide) : ChannelDuplexHandler() {
		private val packetReadBuffer = ByteArrayList()
		private var discardBytes = 0
		private var readableBytes = 0
		private var isCompressed = false
		private var readingType: Type<*>? = null

		private val networkReadBuffer = ByteArrayList()

		override fun isSharable(): Boolean {
			return true
		}

		override fun ensureNotSharable() {}

		private fun drainNetworkBuffer(ctx: ChannelHandlerContext) {
			while (networkReadBuffer.isNotEmpty()) {
				if (discardBytes > 0) {
					val toSkip = discardBytes.coerceAtMost(networkReadBuffer.size)
					discardBytes -= toSkip
					networkReadBuffer.removeElements(0, toSkip)
					continue
				}

				if (readingType != null) {
					while (readableBytes > 0 && networkReadBuffer.isNotEmpty()) {
						val toRead = readableBytes.coerceAtMost(networkReadBuffer.size)
						packetReadBuffer.addElements(packetReadBuffer.size, networkReadBuffer.elements(), 0, toRead)
						readableBytes -= toRead
						networkReadBuffer.removeElements(0, toRead)
					}

					if (readableBytes == 0) {
						val stream: InputStream

						if (isCompressed) {
							stream = BufferedInputStream(LimitingInputStream(InflaterInputStream(FastByteArrayInputStream(packetReadBuffer.elements(), 0, packetReadBuffer.size))))
						} else {
							stream = FastByteArrayInputStream(packetReadBuffer.elements(), 0, packetReadBuffer.size)
						}

						var i = -1

						// legacy protocol allows to stitch multiple packets of same type together without
						// separate headers for each
						// Due to nature of netty pipeline, we can't do the same on native protocol;
						// so don't do that when on native protocol
						do {
							i++

							try {
								ctx.fireChannelRead(readingType!!.factory.read(DataInputStream(stream), isLegacy, side))
							} catch (err: Throwable) {
								LOGGER.error("Error while reading incoming packet from network (type ${readingType!!.id}; ${readingType!!.type}; packet No. $i in stream)", err)
							}

							// are u foking serious, InflaterInputStream?
							if (stream is BufferedInputStream && stream.available() > 0 && isLegacy) {
								stream.mark(1)
								stream.read()
								stream.reset()
							}
						} while (stream.available() > 0 && isLegacy)

						stream.close()

						packetReadBuffer.clear()
						readingType = null
						isCompressed = false
					}

					continue
				}

				val reader = FastByteArrayInputStream(networkReadBuffer.elements(), 0, networkReadBuffer.size)

				try {
					val stream = DataInputStream(reader)
					val packetType = stream.readUnsignedByte()
					val dataLength = stream.readSignedVarInt()

					val type = packets.getOrNull(packetType)

					if (type == null) {
						val name = missingNames[packetType]

						if (name != null)
							LOGGER.error("Unknown packet type $packetType ($name)! Discarding ${dataLength.absoluteValue} bytes")
						else
							LOGGER.error("Unknown packet type $packetType! Discarding ${dataLength.absoluteValue} bytes")

						discardBytes = dataLength.absoluteValue
					} else if (!type.direction.acceptedOn(side)) {
						LOGGER.error("Packet type $packetType (${type.type}) can not be accepted on side $side! Discarding ${dataLength.absoluteValue} bytes")
						discardBytes = dataLength.absoluteValue
					} else if (dataLength.absoluteValue >= MAX_PACKET_SIZE) {
						LOGGER.error("Packet ($packetType/${type.type}) of ${dataLength.absoluteValue} bytes is bigger than maximum allowed $MAX_PACKET_SIZE bytes")
						discardBytes = dataLength.absoluteValue
					} else {
						if (LOG_PACKETS) LOGGER.debug("Packet type {} ({}) received on {} (size {} bytes)", packetType, type.type, side, dataLength.absoluteValue)
						readingType = type
						readableBytes = dataLength.absoluteValue
						isCompressed = dataLength < 0
					}

					networkReadBuffer.removeElements(0, reader.position().toInt())
				} catch (err: EOFException) {
					// Ignore EOF, since it is caused by segmented nature of TCP
					break
				}
			}
		}

		override fun channelRead(ctx: ChannelHandlerContext, msg: Any) {
			if (msg is ByteBuf) {
				try {
					if (msg.readableBytes() > 0) {
						val index = networkReadBuffer.size
						networkReadBuffer.size(index + msg.readableBytes())
						msg.readBytes(networkReadBuffer.elements(), index, msg.readableBytes())
						drainNetworkBuffer(ctx)
					}
				} finally {
					msg.release()
				}
			} else {
				super.channelRead(ctx, msg)
			}
		}

		override fun write(ctx: ChannelHandlerContext, msg: Any, promise: ChannelPromise) {
			val type = clazz2Type[msg::class]

			if (type == null) {
				LOGGER.error("Unknown outgoing message type ${msg::class}, it will not reach the other side.")
			} else if (!type.direction.acceptedOn(side.opposite)) {
				LOGGER.error("Packet ${type.type} can not be accepted on side ${side.opposite}, refusing to send it!")
			} else {
				val stream = FastByteArrayOutputStream()
				(msg as IPacket).write(DataOutputStream(stream), isLegacy)

				if (isLegacy && stream.length == 0)
					throw IllegalStateException("Packet $msg didn't write any data to network, this is not allowed by legacy protocol")

				if (stream.length >= 512) {
					// compress
					val deflater = Deflater(3)
					val buffers = ByteArrayList(4096)
					val buffer = ByteArray(4096)
					deflater.setInput(stream.array, 0, stream.length)
					deflater.finish()

					while (!deflater.finished()) {
						val deflated = deflater.deflate(buffer)

						if (deflated > 0)
							buffers.addElements(buffers.size, buffer, 0, deflated)
						else
							break
					}

					val buff = ctx.alloc().buffer(buffers.size + 5)
					val stream2 = ByteBufOutputStream(buff)
					stream2.writeByte(type.id)
					stream2.writeSignedVarInt(-buffers.size)
					stream2.write(buffers.elements(), 0, buffers.size)
					if (LOG_PACKETS) LOGGER.debug("Packet type {} ({}) sent from {} (size {} bytes / COMPRESSED size {} bytes)", type.id, type.type, side, stream.length, buffers.size)
					ctx.write(buff, promise)
				} else {
					// send as-is
					val buff = ctx.alloc().buffer(stream.length + 5)
					val stream2 = ByteBufOutputStream(buff)
					stream2.writeByte(type.id)
					stream2.writeSignedVarInt(stream.length)
					stream2.write(stream.array, 0, stream.length)
					if (LOG_PACKETS) LOGGER.debug("Packet type {} ({}) sent from {} (size {} bytes)", type.id, type.type, side, stream.length)
					ctx.write(buff, promise)
				}
			}
		}
	}

	inner class Validator(val side: ConnectionSide) : ChannelDuplexHandler() {
		override fun channelRead(ctx: ChannelHandlerContext, msg: Any) {
			val type = clazz2Type[msg::class]

			if (type == null) {
				LOGGER.error("Unknown packet type ${msg::class}!")
			} else if (!type.direction.acceptedOn(side)) {
				LOGGER.error("Packet ${type.type} can not be accepted on side $side!")
			} else {
				try {
					ctx.fireChannelRead(msg)
				} catch (err: Throwable) {
					LOGGER.error("Error while reading incoming packet from network", err)
				}
			}
		}

		override fun write(ctx: ChannelHandlerContext, msg: Any, promise: ChannelPromise) {
			val type = clazz2Type[msg::class]

			if (type == null) {
				LOGGER.error("Unknown outgoing message type ${msg::class}, it will not reach the other side.")
			} else if (!type.direction.acceptedOn(side.opposite)) {
				LOGGER.error("Packet ${type.type} can not be accepted on side ${side.opposite}, refusing to send it!")
			} else {
				ctx.write(msg, promise)
			}
		}
	}

	companion object {
		const val LOG_PACKETS = false
		const val MAX_PACKET_SIZE = 64L * 1024L * 1024L // 64 MiB
		// this includes both compressed and uncompressed
		// Original game allows 16 mebibyte packets
		// but it doesn't account for compression bomb (packets are fully uncompressed
		// right away without limiting decompressed output size)

		private val LOGGER = LogManager.getLogger()

		val NATIVE = PacketRegistry(false)
		val LEGACY = PacketRegistry(true)
		val HANDSHAKE = PacketRegistry(false)

		init {
			NATIVE.add(::JoinWorldPacket)
			NATIVE.add(::ChunkCellsPacket)
			NATIVE.add(::ForgetChunkPacket)
			NATIVE.add(::SpawnWorldObjectPacket)
			NATIVE.add(::ForgetEntityPacket)
			NATIVE.add(::UniverseTimeUpdatePacket)
			NATIVE.add(::StepUpdatePacket)

			HANDSHAKE.add(::ProtocolRequestPacket)
			HANDSHAKE.add(::ProtocolResponsePacket)

			// legacy protocol handshake looks like this
			// --> ProtocolRequest
			// <-- ProtocolResponse
			// --> ClientConnect
			// <-- HandshakeChallenge *
			// --> HandshakeResponse *
			// <-- ConnectSuccess / ConnectFailure
			// <-- UniverseClockUpdatePacket
			// <-- WorldStartPacket

			LEGACY.skip("ProtocolRequest")
			LEGACY.skip("ProtocolResponse")

			// Packets sent universe server -> universe client
			LEGACY.add(::ServerDisconnectPacket) // ServerDisconnect
			LEGACY.add(::ConnectSuccessPacket) // ConnectSuccess
			LEGACY.add(::ConnectFailurePacket)
			LEGACY.add(::HandshakeChallengePacket) // HandshakeChallenge
			LEGACY.add(::ChatReceivePacket)
			LEGACY.add(::UniverseTimeUpdatePacket)
			LEGACY.add(::CelestialResponsePacket)
			LEGACY.add(::PlayerWarpResultPacket)
			LEGACY.skip("PlanetTypeUpdate")
			LEGACY.skip("Pause")
			LEGACY.add(::ServerInfoPacket)

			// Packets sent universe client -> universe server
			LEGACY.add(::ClientConnectPacket) // ClientConnect
			LEGACY.add(ClientDisconnectRequestPacket::read)
			LEGACY.add(::HandshakeResponsePacket) // HandshakeResponse
			LEGACY.add(::PlayerWarpPacket)
			LEGACY.add(::FlyShipPacket)
			LEGACY.add(::ChatSendPacket)
			LEGACY.add(::CelestialRequestPacket)

			// Packets sent bidirectionally between the universe client and the universe
			// server
			LEGACY.add(ClientContextUpdatePacket::read)

			// Packets sent world server -> world client
			LEGACY.add(::WorldStartPacket) // WorldStart
			LEGACY.add(::WorldStopPacket)
			LEGACY.skip("WorldLayoutUpdate")
			LEGACY.skip("WorldParametersUpdate")
			LEGACY.add(::CentralStructureUpdatePacket)
			LEGACY.add(LegacyTileArrayUpdatePacket::read)
			LEGACY.add(LegacyTileUpdatePacket::read)
			LEGACY.skip("TileLiquidUpdate")
			LEGACY.add(::TileDamageUpdatePacket)
			LEGACY.add(::TileModificationFailurePacket)
			LEGACY.add(::GiveItemPacket)
			LEGACY.add(::EnvironmentUpdatePacket)
			LEGACY.skip("UpdateTileProtection")
			LEGACY.skip("SetDungeonGravity")
			LEGACY.skip("SetDungeonBreathable")
			LEGACY.add(::SetPlayerStartPacket)
			LEGACY.add(FindUniqueEntityResponsePacket::read)
			LEGACY.add(PongPacket::read)

			// Packets sent world client -> world server
			LEGACY.add(::ModifyTileListPacket)
			LEGACY.add(::DamageTileGroupPacket)
			LEGACY.skip("CollectLiquid")
			LEGACY.add(::RequestDropPacket)
			LEGACY.add(::SpawnEntityPacket)
			LEGACY.add(::ConnectWirePacket)
			LEGACY.skip("DisconnectAllWires")
			LEGACY.add(::WorldClientStateUpdatePacket)
			LEGACY.add(::FindUniqueEntityPacket)
			LEGACY.add(WorldStartAcknowledgePacket::read)
			LEGACY.add(PingPacket::read)

			// Packets sent bidirectionally between world client and world server
			LEGACY.add(::EntityCreatePacket)
			LEGACY.add(EntityUpdateSetPacket::read)
			LEGACY.add(::EntityDestroyPacket)
			LEGACY.add(::EntityInteractPacket)
			LEGACY.add(::EntityInteractResultPacket)
			LEGACY.skip("HitRequest")
			LEGACY.skip("DamageRequest")
			LEGACY.skip("DamageNotification")
			LEGACY.skip("EntityMessage")
			LEGACY.skip("EntityMessageResponse")
			LEGACY.add(::UpdateWorldPropertiesPacket)
			LEGACY.add(::StepUpdatePacket)

			// Packets sent system server -> system client
			LEGACY.add(::SystemWorldStartPacket)
			LEGACY.add(::SystemWorldUpdatePacket)
			LEGACY.add(::SystemObjectCreatePacket)
			LEGACY.add(::SystemObjectDestroyPacket)
			LEGACY.add(::SystemShipCreatePacket)
			LEGACY.add(::SystemShipDestroyPacket)

			// Packets sent system client -> system server
			LEGACY.skip("SystemObjectSpawn")
		}
	}
}
