blade-addons-next

package blade.addon.utils.rendering;

import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import net.fabricmc.fabric.api.client.rendering.v1.level.LevelRenderContext;
import net.minecraft.client.renderer.rendertype.RenderType;
import net.minecraft.client.renderer.state.level.LevelRenderState;
import net.minecraft.world.phys.Vec3;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

public class RenderHandler {

    private final RenderType renderType;

    public RenderHandler(RenderType renderType)  {
        this.renderType = renderType;
    }

    private final List<RenderingEvent> listeners = new ArrayList<>();

    public void register(RenderingEvent listener) {
        listeners.add(listener);
    }

    public void invoke(Consumer<RenderingEvent> action) {
        if (listeners.isEmpty()) return;
        for (RenderingEvent listener : listeners) {
            action.accept(listener);
        }
    }

    public void init(LevelRenderContext context) {
        LevelRenderState worldState = context.levelState();
        if (worldState == null) return;
        Vec3 camera = worldState.cameraRenderState.pos;
        PoseStack matrices = context.poseStack();
        if (matrices == null) return;
        if (listeners.isEmpty()) return;
        matrices.pushPose();
        matrices.translate(-camera.x, -camera.y, -camera.z);

        // Listeners run during the submit-collection phase (they also submit Gizmos and text, which
        // are only accepted while collecting). Their vertex output is recorded and replayed into the
        // real per-frame consumer when the custom geometry submit is drawn.
        RecordingVertexConsumer recording = new RecordingVertexConsumer();
        invoke(renderingEvent -> renderingEvent.render(context, matrices, recording));

        context.submitNodeCollector().submitCustomGeometry(matrices, renderType, (pose, consumer) -> recording.replay(consumer));
        matrices.popPose();
    }

    private static final class RecordingVertexConsumer implements VertexConsumer {
        private final List<Consumer<VertexConsumer>> ops = new ArrayList<>();

        @Override
        public VertexConsumer addVertex(float x, float y, float z) {
            ops.add(consumer -> consumer.addVertex(x, y, z));
            return this;
        }

        @Override
        public VertexConsumer setColor(int r, int g, int b, int a) {
            ops.add(consumer -> consumer.setColor(r, g, b, a));
            return this;
        }

        @Override
        public VertexConsumer setColor(int rgba) {
            ops.add(consumer -> consumer.setColor(rgba));
            return this;
        }

        @Override
        public VertexConsumer setUv(float u, float v) {
            ops.add(consumer -> consumer.setUv(u, v));
            return this;
        }

        @Override
        public VertexConsumer setUv1(int u, int v) {
            ops.add(consumer -> consumer.setUv1(u, v));
            return this;
        }

        @Override
        public VertexConsumer setUv2(int u, int v) {
            ops.add(consumer -> consumer.setUv2(u, v));
            return this;
        }

        // New abstract method on VertexConsumer in 26.3 (third UV set for the renderpearl vertex formats).
        @Override
        public VertexConsumer setUv3(float u, float v) {
            ops.add(consumer -> consumer.setUv3(u, v));
            return this;
        }

        @Override
        public VertexConsumer setNormal(float x, float y, float z) {
            ops.add(consumer -> consumer.setNormal(x, y, z));
            return this;
        }

        @Override
        public VertexConsumer setLineWidth(float width) {
            ops.add(consumer -> consumer.setLineWidth(width));
            return this;
        }

        private void replay(VertexConsumer consumer) {
            for (Consumer<VertexConsumer> op : ops) {
                op.accept(consumer);
            }
        }
    }

}