smog/shared/src/main/java/me/lethunderhawk/network/packet/PacketDispatcher.java

package me.lethunderhawk.network.packet;

import com.esotericsoftware.kryonet.Connection;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

public class PacketDispatcher {

    private static final Map<Class<? extends Packet>, List<RegisteredListener>> listeners = new ConcurrentHashMap<>();

    public static void registerListener(PacketListener listener) {
        for (Method method : listener.getClass().getDeclaredMethods()) {
            if (!method.isAnnotationPresent(PacketHandler.class)) {
                continue;
            }

            Class<?>[] params = method.getParameterTypes();
            if (params.length != 1 && params.length != 2) {
                continue;
            }

            if (!Packet.class.isAssignableFrom(params[0])) {
                continue;
            }
            if (params.length == 2 && !Connection.class.isAssignableFrom(params[1])) {
                continue;
            }

            method.setAccessible(true);

            listeners.computeIfAbsent((Class<? extends Packet>) params[0], c -> new CopyOnWriteArrayList<>())
                .add(new RegisteredListener(listener, method, params.length));
        }
    }

    public static void dispatchPacket(Packet packet, Connection connection) {
        List<RegisteredListener> packetListeners = listeners.get(packet.getClass());
        if (packetListeners == null || packetListeners.isEmpty()) {
            return;
        }

        for (RegisteredListener entry : packetListeners) {
            try {
                if (entry.argCount == 2) {
                    entry.method.invoke(entry.listener, packet, connection);
                } else {
                    entry.method.invoke(entry.listener, packet);
                }
            } catch (InvocationTargetException e) {
                // Unwrap the actual listener exception for clearer stack traces
                throw new RuntimeException("Packet listener failed for " + packet.getClass().getSimpleName(),
                    e.getCause() != null ? e.getCause() : e);
            } catch (Exception e) {
                throw new RuntimeException("Failed to invoke listener for " + packet.getClass().getSimpleName(), e);
            }
        }
    }

    public static void dispatchPacket(Packet packet) {
        dispatchPacket(packet, null);
    }

    private static class RegisteredListener {
        private final PacketListener listener;
        private final Method method;
        private final int argCount;

        public RegisteredListener(PacketListener listener, Method method, int argCount) {
            this.listener = listener;
            this.method = method;
            this.argCount = argCount;
        }
    }
}