LibJF/libjf-config-network-v0/src/main/java/io/gitlab/jfronny/libjf/config/impl/network/RequestRouter.java

package io.gitlab.jfronny.libjf.config.impl.network;

import io.gitlab.jfronny.libjf.LibJf;
import net.fabricmc.fabric.api.networking.v1.PacketByteBufs;
import net.fabricmc.fabric.api.networking.v1.PacketSender;
import net.minecraft.network.PacketByteBuf;
import net.minecraft.util.Identifier;

import java.util.*;

public class RequestRouter {
    public static final String MOD_ID = "libjf-config-network-v0";
    public static final Identifier RESPONSE_ID = new Identifier(MOD_ID, "response");
    public static final Identifier REQUEST_ID = new Identifier(MOD_ID, "request");
    public static int PROTOCOL_VERSION = 1;

    private static final Map<String, RequestHandler> persistendHandlers = new HashMap<>();
    private static final Map<Long, Request> currentRequests = new HashMap<>(); //TODO implement timeout and prune old requests
    private static final Random random = new Random();

    public static void acceptResponse(PacketByteBuf buf, PacketSender responseSender) {
        Request request = currentRequests.remove(buf.readLong());
        if (request != null) {
            switch (buf.readInt()) {
                case 0 -> request.responseHandler.onSuccess(buf);
                case 1 -> request.responseHandler.onNotFound();
                case 2 -> request.responseHandler.onDeny();
                case 3 -> request.responseHandler.onFailure(buf.readString());
                default -> request.responseHandler.onFailure("Unrecognized error received");
            }
        }
    }

    public static void acceptRequest(PacketByteBuf buf, PacketSender responseSender) {
        PacketByteBuf resp = PacketByteBufs.create();
        long id = buf.readLong();
        resp.writeLong(id);
        try {
            if (buf.readBoolean()) {
                // persistent
                handleRequest(resp, id, buf, persistendHandlers.get(buf.readString()));
            } else {
                // followup
                Request parent = currentRequests.get(buf.readLong());
                if (parent == null) {
                    resp.writeInt(1);
                } else {
                    String key = buf.readString();
                    RequestHandler handler = parent.temporaryHandlers.get(key);
                    if (handler == null) handler = persistendHandlers.get(key);
                    handleRequest(resp, id, buf, handler);
                }
            }
        } catch (Throwable t) {
            LibJf.LOGGER.error("Cannot complete request", t);
            resp.writeInt(3);
            resp.writeString(t.getMessage() == null ? "null" : t.getMessage());
        } finally {
            responseSender.sendPacket(RESPONSE_ID, resp);
        }
    }

    private static void handleRequest(PacketByteBuf resp, long id, PacketByteBuf buf, RequestHandler handler) throws Throwable {
        if (handler == null) {
            resp.writeInt(1);
        } else {
            PacketByteBuf response = handler.handle(
                    buf,
                    (responseSender, name, body, responseHandler, temporaryHandlers) ->
                            RequestRouter.sendRequest(responseSender, id, name, body, responseHandler, temporaryHandlers)
            );
            resp.writeInt(0);
            if (response != null) resp.writeBytes(response.copy());
        }
    }

    public static void deny(PacketByteBuf buf, PacketSender responseSender) {
        PacketByteBuf resp = PacketByteBufs.create();
        resp.writeLong(buf.readLong());
        resp.writeInt(2);
        responseSender.sendPacket(RESPONSE_ID, resp);
    }

    public static void registerHandler(String name, RequestHandler handler) {
        persistendHandlers.put(name, handler);
    }

    public static void sendRequest(PacketSender responseSender, String name, PacketByteBuf body, ResponseHandler responseHandler, Map<String, RequestHandler> temporaryHandlers) {
        sendRequest(responseSender, null, name, body, responseHandler, temporaryHandlers);
    }

    private static void sendRequest(PacketSender responseSender, Long parent, String name, PacketByteBuf body, ResponseHandler responseHandler, Map<String, RequestHandler> temporaryHandlers) {
        long id;
        synchronized (currentRequests) {
            Set<Long> keys = currentRequests.keySet();
            do {
                id = random.nextLong();
            } while (keys.contains(id));
            currentRequests.put(id, new Request(temporaryHandlers, responseHandler));
        }
        PacketByteBuf req = PacketByteBufs.create();
        req.writeLong(id);
        req.writeBoolean(parent == null);
        if (parent != null) req.writeLong(parent);
        req.writeString(name);
        if (body != null) {
            req.writeBytes(body.copy());
        }
        responseSender.sendPacket(REQUEST_ID, req);
    }

    private record Request(Map<String, RequestHandler> temporaryHandlers, ResponseHandler responseHandler) {}
}