package ru.dbotthepony.mc.otm.network

import com.mojang.serialization.Codec
import io.netty.buffer.ByteBuf
import io.netty.buffer.ByteBufAllocator
import it.unimi.dsi.fastutil.bytes.ByteArrayList
import net.minecraft.core.RegistryAccess
import net.minecraft.network.FriendlyByteBuf
import net.minecraft.network.RegistryFriendlyByteBuf
import net.minecraft.network.codec.ByteBufCodecs
import net.minecraft.network.codec.StreamCodec
import net.minecraft.network.protocol.common.custom.CustomPacketPayload
import net.neoforged.neoforge.network.connection.ConnectionType
import net.neoforged.neoforge.network.handling.IPayloadContext
import net.neoforged.neoforge.network.registration.PayloadRegistrar
import java.util.Optional
import kotlin.reflect.KFunction1

fun <T : CustomPacketPayload> PayloadRegistrar.playToClient(
	type: CustomPacketPayload.Type<T>,
	codec: StreamCodec<FriendlyByteBuf, T>,
	handler: KFunction1<IPayloadContext, Unit>
): PayloadRegistrar = playToClient(type, codec) { _, context -> handler(context) }

fun <T : CustomPacketPayload> PayloadRegistrar.playToServer(
	type: CustomPacketPayload.Type<T>,
	codec: StreamCodec<FriendlyByteBuf, T>,
	handler: KFunction1<IPayloadContext, Unit>
): PayloadRegistrar = playToServer(type, codec) { _, context -> handler(context) }

inline fun encodePayload(registry: RegistryAccess, block: (RegistryFriendlyByteBuf) -> Unit): ByteArrayList {
	val underlying = ByteBufAllocator.DEFAULT.buffer()
	val buf = RegistryFriendlyByteBuf(underlying, registry, ConnectionType.NEOFORGE)

	try {
		block.invoke(buf)
		buf.readerIndex(0)

		val bytes = ByteArray(buf.readableBytes())
		buf.readBytes(bytes)
		return ByteArrayList.wrap(bytes)
	} finally {
		underlying.release()
	}
}

inline fun <T> decodePayload(registry: RegistryAccess, data: ByteArrayList, block: (RegistryFriendlyByteBuf) -> T): T {
	val underlying = ByteBufAllocator.DEFAULT.buffer(data.size)
	val buf = RegistryFriendlyByteBuf(underlying, registry, ConnectionType.NEOFORGE)

	try {
		underlying.writeBytes(data.elements(), 0, data.size)
		return block.invoke(buf)
	} finally {
		underlying.release()
	}
}

fun FriendlyByteBuf.writeByteListUnbounded(list: ByteArrayList) {
	writeBytes(list.elements(), 0, list.size)
}

fun FriendlyByteBuf.readByteListUnbounded(): ByteArrayList {
	return ByteArrayList.wrap(ByteArray(readableBytes()).also(::readBytes))
}

fun FriendlyByteBuf.writeByteList(list: ByteArrayList) {
	writeVarInt(list.size)
	writeBytes(list.elements(), 0, list.size)
}

fun FriendlyByteBuf.readByteList(): ByteArrayList {
	return ByteArrayList.wrap(ByteArray(readVarInt()).also(::readBytes))
}

fun <S : ByteBuf, V> StreamCodec<S, V>.wrap(): MatteryStreamCodec<S, V> = MatteryStreamCodec.Wrapper(this)
@Deprecated("Redundant wrapping", level = DeprecationLevel.ERROR, replaceWith = ReplaceWith("this"))
fun <S : ByteBuf, V> MatteryStreamCodec<S, V>.wrap(): MatteryStreamCodec<S, V> = this
fun <S : ByteBuf, V> MatteryStreamCodec<S, V>.nullable(): MatteryStreamCodec<S, V?> = MatteryStreamCodec.Nullable(this)
fun <S : ByteBuf, V : Any> MatteryStreamCodec<S, V>.optional(): MatteryStreamCodec<S, Optional<V>> = MatteryStreamCodec.Optional(this)

fun <V> StreamCodec<in RegistryFriendlyByteBuf, V>.encode(registry: RegistryAccess, value: V): ByteArrayList {
	return encodePayload(registry) {
		encode(it, value)
	}
}

fun <V> StreamCodec<in RegistryFriendlyByteBuf, V>.decode(registry: RegistryAccess, value: ByteArrayList): V {
	return decodePayload(registry, value) {
		decode(it)
	}
}

fun <T> Codec<T>.streamCodec(): MatteryStreamCodec<RegistryFriendlyByteBuf, T> {
	return ByteBufCodecs.fromCodecWithRegistries(this).wrap()
}
