MyDimension-Lord

package com.xfestudio.mydimensionlord.platform;

import net.minecraft.network.FriendlyByteBuf;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.level.ServerPlayer;
import net.minecraftforge.network.ChannelBuilder;
import net.minecraftforge.network.PacketDistributor;
import net.minecraftforge.network.SimpleChannel;
import net.minecraftforge.network.NetworkDirection;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Supplier;

public final class PacketChannel {
    private final SimpleChannel nativeChannel;
    public PacketChannel(ResourceLocation name, String protocol) {
        nativeChannel = ChannelBuilder.named(name).networkProtocolVersion(Integer.parseInt(protocol)).simpleChannel();
    }
    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); }
    public void sendToServer(Object packet) { nativeChannel.send(packet, PacketDistributor.SERVER.noArg()); }
    public void sendToPlayer(ServerPlayer player, Object packet) { nativeChannel.send(packet, PacketDistributor.PLAYER.with(player)); }

    public final class Message<T> {
        private final Class<T> type;
        private final int id;
        private final PacketDirection direction;
        private BiConsumer<T, FriendlyByteBuf> encoder;
        private Function<FriendlyByteBuf, T> decoder;
        private BiConsumer<T, Supplier<PacketContext>> handler;
        private Message(Class<T> type, int id, PacketDirection direction) { this.type=type; this.id=id; this.direction=direction; }
        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() {
            nativeChannel.messageBuilder(type, id, direction == PacketDirection.PLAY_TO_SERVER ? NetworkDirection.PLAY_TO_SERVER : NetworkDirection.PLAY_TO_CLIENT)
                    .encoder((packet, buffer) -> encoder.accept(packet, buffer)).decoder(buffer -> decoder.apply(buffer))
                    .consumerNetworkThread((packet, context) -> {
                        PacketContext adapter = new PacketContext(context.getSender(), context::enqueueWork);
                        handler.accept(packet, () -> adapter);
                        context.setPacketHandled(true);
                    }).add();
        }
    }
}