MyDimension-Lord

package com.xfestudio.mydimensionlord.platform;

import com.xfestudio.mydimensionlord.MyDimension;
import net.minecraft.network.FriendlyByteBuf;
import net.minecraft.network.RegistryFriendlyByteBuf;
import net.minecraft.network.codec.StreamCodec;
import net.minecraft.network.protocol.common.custom.CustomPacketPayload;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.level.ServerPlayer;
import net.neoforged.bus.api.SubscribeEvent;
import net.neoforged.fml.common.EventBusSubscriber;
import net.neoforged.neoforge.network.PacketDistributor;
import net.neoforged.neoforge.network.event.RegisterPayloadHandlersEvent;
import net.neoforged.neoforge.network.registration.HandlerThread;
import net.neoforged.neoforge.network.registration.PayloadRegistrar;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Supplier;

@EventBusSubscriber(modid = MyDimension.MOD_ID, bus = EventBusSubscriber.Bus.MOD)
public final class PacketChannel {
    private static PacketChannel instance;
    private final ResourceLocation name;
    private final String protocol;
    private final Map<Class<?>, Message<?>> messages = new LinkedHashMap<>();

    public PacketChannel(ResourceLocation name, String protocol) {
        this.name = name;
        this.protocol = protocol;
        instance = this;
    }

    @SubscribeEvent
    public static void registerPayloads(RegisterPayloadHandlersEvent event) {
        if (instance == null) throw new IllegalStateException("Packets must be initialized during mod construction");
        PayloadRegistrar registrar = event.registrar(instance.protocol).executesOn(HandlerThread.NETWORK);
        instance.messages.values().forEach(message -> message.register(registrar));
    }

    public <T> void registerMessage(int id, Class<T> type, BiConsumer<T, FriendlyByteBuf> encoder,
                                   Function<FriendlyByteBuf, T> decoder, BiConsumer<T, Supplier<PacketContext>> handler) {
        messageBuilder(type, id, id < 2 ? PacketDirection.PLAY_TO_SERVER : PacketDirection.PLAY_TO_CLIENT)
                .encoder(encoder).decoder(decoder).consumerNetworkThread(handler).add();
    }

    public <T> Message<T> messageBuilder(Class<T> type, int id, PacketDirection direction) {
        return new Message<>(type, id, direction);
    }

    @SuppressWarnings("unchecked")
    private <T> Payload<T> wrap(T packet, PacketDirection expected) {
        Message<T> message = (Message<T>) messages.get(packet.getClass());
        if (message == null || message.direction != expected) throw new IllegalArgumentException("Wrong packet direction");
        return new Payload<>(message, packet);
    }

    public void sendToServer(Object packet) {
        PacketDistributor.sendToServer(wrap(packet, PacketDirection.PLAY_TO_SERVER));
    }

    public void sendToPlayer(ServerPlayer player, Object packet) {
        PacketDistributor.sendToPlayer(player, wrap(packet, PacketDirection.PLAY_TO_CLIENT));
    }

    private record Payload<T>(Message<T> message, T value) implements CustomPacketPayload {
        @Override public Type<? extends CustomPacketPayload> type() { return message.payloadType; }
    }

    public final class Message<T> {
        private final Class<T> javaType;
        private final PacketDirection direction;
        private final CustomPacketPayload.Type<Payload<T>> payloadType;
        private BiConsumer<T, FriendlyByteBuf> encoder;
        private Function<FriendlyByteBuf, T> decoder;
        private BiConsumer<T, Supplier<PacketContext>> handler;

        private Message(Class<T> javaType, int id, PacketDirection direction) {
            this.javaType = javaType;
            this.direction = direction;
            payloadType = new CustomPacketPayload.Type<>(ResourceLocation.fromNamespaceAndPath(
                    name.getNamespace(), name.getPath() + "/" + id));
        }

        public Message<T> encoder(BiConsumer<T, FriendlyByteBuf> encoder) { this.encoder = encoder; return this; }
        public Message<T> decoder(Function<FriendlyByteBuf, T> decoder) { this.decoder = decoder; return this; }
        public Message<T> consumerNetworkThread(BiConsumer<T, Supplier<PacketContext>> handler) { this.handler = handler; return this; }
        public Message<T> consumerMainThread(BiConsumer<T, Supplier<PacketContext>> handler) {
            this.handler = (packet, context) -> context.get().enqueueWork(() -> handler.accept(packet, context));
            return this;
        }
        public void add() {
            if (messages.putIfAbsent(javaType, this) != null) throw new IllegalStateException("Duplicate packet " + javaType);
        }
        private void register(PayloadRegistrar registrar) {
            StreamCodec<RegistryFriendlyByteBuf, Payload<T>> codec = new StreamCodec<>() {
                @Override public Payload<T> decode(RegistryFriendlyByteBuf buf) { return new Payload<>(Message.this, decoder.apply(buf)); }
                @Override public void encode(RegistryFriendlyByteBuf buf, Payload<T> packet) { encoder.accept(packet.value, buf); }
            };
            net.neoforged.neoforge.network.handling.IPayloadHandler<Payload<T>> receiver = (packet, context) -> {
                // Clientbound play packets can arrive before the local player has
                // been created. Only serverbound handlers need a sender.
                ServerPlayer sender = direction == PacketDirection.PLAY_TO_SERVER
                        ? (ServerPlayer) context.player() : null;
                PacketContext adapter = new PacketContext(sender, context::enqueueWork);
                handler.accept(packet.value, () -> adapter);
            };
            if (direction == PacketDirection.PLAY_TO_SERVER) registrar.playToServer(payloadType, codec, receiver);
            else registrar.playToClient(payloadType, codec, receiver);
        }
    }
}