TiedUp-/src/main/java/com/tiedup/remake/client/gltf/GltfMeshRenderer.java

package com.tiedup.remake.client.gltf;

import com.mojang.blaze3d.vertex.DefaultVertexFormat;
import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import com.mojang.blaze3d.vertex.VertexFormat;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import net.minecraft.client.renderer.MultiBufferSource;
import net.minecraft.client.renderer.RenderStateShard;
import net.minecraft.client.renderer.RenderType;
import net.minecraft.resources.ResourceLocation;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import org.joml.Matrix3f;
import org.joml.Matrix4f;
import org.joml.Vector4f;

/**
 * Submits CPU-skinned glTF mesh vertices to Minecraft's rendering pipeline.
 * Uses TRIANGLES mode RenderType (same pattern as ObjModelRenderer).
 */
@OnlyIn(Dist.CLIENT)
public final class GltfMeshRenderer extends RenderStateShard {

    private static final ResourceLocation WHITE_TEXTURE =
        ResourceLocation.fromNamespaceAndPath(
            "tiedup",
            "models/obj/shared/white.png"
        );

    /** Cached default RenderType (white texture). Created once, reused every frame. */
    private static RenderType cachedDefaultRenderType;

    /** Cache for texture-specific RenderTypes, keyed by ResourceLocation. */
    private static final Map<ResourceLocation, RenderType> RENDER_TYPE_CACHE =
        new ConcurrentHashMap<>();

    private GltfMeshRenderer() {
        super("tiedup_gltf_renderer", () -> {}, () -> {});
    }

    /**
     * Get the default TRIANGLES-mode RenderType (white texture), creating it once if needed.
     */
    private static RenderType getDefaultRenderType() {
        if (cachedDefaultRenderType == null) {
            cachedDefaultRenderType = createTriangleRenderType(WHITE_TEXTURE);
        }
        return cachedDefaultRenderType;
    }

    /**
     * Public accessor for the default RenderType (white texture).
     * Used by external renderers that need the same RenderType for tinted rendering.
     */
    public static RenderType getRenderTypeForDefaultTexture() {
        return getDefaultRenderType();
    }

    /**
     * Get a RenderType for a specific texture, caching it for reuse.
     *
     * @param texture the texture ResourceLocation
     * @return the cached or newly created RenderType
     */
    private static RenderType getRenderTypeForTexture(
        ResourceLocation texture
    ) {
        return RENDER_TYPE_CACHE.computeIfAbsent(
            texture,
            GltfMeshRenderer::createTriangleRenderType
        );
    }

    /**
     * Create a TRIANGLES-mode RenderType for glTF mesh rendering with the given texture.
     */
    private static RenderType createTriangleRenderType(
        ResourceLocation texture
    ) {
        RenderType.CompositeState state = RenderType.CompositeState.builder()
            .setShaderState(RENDERTYPE_ENTITY_CUTOUT_NO_CULL_SHADER)
            .setTextureState(
                new RenderStateShard.TextureStateShard(texture, false, false)
            )
            .setTransparencyState(NO_TRANSPARENCY)
            .setCullState(NO_CULL)
            .setLightmapState(LIGHTMAP)
            .setOverlayState(OVERLAY)
            .createCompositeState(true);

        return RenderType.create(
            "tiedup_gltf_triangles",
            DefaultVertexFormat.NEW_ENTITY,
            VertexFormat.Mode.TRIANGLES,
            256 * 1024,
            true,
            false,
            state
        );
    }

    /**
     * Clear cached RenderTypes. Call on resource reload so that re-exported
     * textures are picked up without restarting the game.
     */
    public static void clearRenderTypeCache() {
        cachedDefaultRenderType = null;
        RENDER_TYPE_CACHE.clear();
    }

    /**
     * Render a skinned glTF mesh using the default white texture.
     *
     * @param data          parsed glTF data
     * @param jointMatrices computed joint matrices from skinning engine
     * @param poseStack     current pose stack
     * @param buffer        multi-buffer source
     * @param packedLight   packed light value
     * @param packedOverlay packed overlay value
     */
    public static void renderSkinned(
        GltfData data,
        Matrix4f[] jointMatrices,
        PoseStack poseStack,
        MultiBufferSource buffer,
        int packedLight,
        int packedOverlay
    ) {
        renderSkinnedInternal(
            data,
            jointMatrices,
            poseStack,
            buffer,
            packedLight,
            packedOverlay,
            getDefaultRenderType()
        );
    }

    /**
     * Render a skinned glTF mesh using a custom texture.
     *
     * @param data          parsed glTF data
     * @param jointMatrices computed joint matrices from skinning engine
     * @param poseStack     current pose stack
     * @param buffer        multi-buffer source
     * @param packedLight   packed light value
     * @param packedOverlay packed overlay value
     * @param texture       the texture to use for rendering
     */
    public static void renderSkinned(
        GltfData data,
        Matrix4f[] jointMatrices,
        PoseStack poseStack,
        MultiBufferSource buffer,
        int packedLight,
        int packedOverlay,
        ResourceLocation texture
    ) {
        renderSkinnedInternal(
            data,
            jointMatrices,
            poseStack,
            buffer,
            packedLight,
            packedOverlay,
            getRenderTypeForTexture(texture)
        );
    }

    /**
     * Internal rendering implementation shared by both overloads.
     */
    private static void renderSkinnedInternal(
        GltfData data,
        Matrix4f[] jointMatrices,
        PoseStack poseStack,
        MultiBufferSource buffer,
        int packedLight,
        int packedOverlay,
        RenderType renderType
    ) {
        Matrix4f pose = poseStack.last().pose();
        Matrix3f normalMat = poseStack.last().normal();

        VertexConsumer vc = buffer.getBuffer(renderType);

        int[] indices = data.indices();
        float[] texCoords = data.texCoords();

        float[] outPos = new float[3];
        float[] outNormal = new float[3];

        // Pre-allocate scratch vectors outside the loop to avoid per-vertex allocations
        Vector4f tmpPos = new Vector4f();
        Vector4f tmpNorm = new Vector4f();

        for (int idx : indices) {
            // Skin this vertex
            GltfSkinningEngine.skinVertex(
                data,
                idx,
                jointMatrices,
                outPos,
                outNormal,
                tmpPos,
                tmpNorm
            );

            // UV coordinates
            float u = texCoords[idx * 2];
            float v = texCoords[idx * 2 + 1];

            vc
                .vertex(pose, outPos[0], outPos[1], outPos[2])
                .color(255, 255, 255, 255)
                .uv(u, 1.0f - v)
                .overlayCoords(packedOverlay)
                .uv2(packedLight)
                .normal(normalMat, outNormal[0], outNormal[1], outNormal[2])
                .endVertex();
        }
    }

    /**
     * Render a skinned glTF mesh with per-primitive tint colors.
     *
     * <p>Each primitive in the mesh is checked against the tintColors map.
     * If a primitive is tintable and its channel is present in the map,
     * the corresponding RGB color is applied as vertex color (multiplied
     * against the texture by the {@code rendertype_entity_cutout_no_cull} shader).
     * Non-tintable primitives render with white (no tint).</p>
     *
     * <p>This is a single VertexConsumer stream — all primitives share the
     * same RenderType and draw call, only the vertex color differs per range.</p>
     *
     * @param data          parsed glTF data (must have primitives)
     * @param jointMatrices computed joint matrices from skinning engine
     * @param poseStack     current pose stack
     * @param buffer        multi-buffer source
     * @param packedLight   packed light value
     * @param packedOverlay packed overlay value
     * @param renderType    the RenderType to use
     * @param tintColors    channel name to RGB int (0xRRGGBB); empty map = white everywhere
     */
    public static void renderSkinnedTinted(
        GltfData data,
        Matrix4f[] jointMatrices,
        PoseStack poseStack,
        MultiBufferSource buffer,
        int packedLight,
        int packedOverlay,
        RenderType renderType,
        Map<String, Integer> tintColors
    ) {
        Matrix4f pose = poseStack.last().pose();
        Matrix3f normalMat = poseStack.last().normal();

        VertexConsumer vc = buffer.getBuffer(renderType);
        float[] texCoords = data.texCoords();

        float[] outPos = new float[3];
        float[] outNormal = new float[3];
        Vector4f tmpPos = new Vector4f();
        Vector4f tmpNorm = new Vector4f();

        List<GltfData.Primitive> primitives = data.primitives();

        for (GltfData.Primitive prim : primitives) {
            // Determine color for this primitive
            int r = 255,
                g = 255,
                b = 255;
            if (prim.tintable() && prim.tintChannel() != null) {
                Integer colorInt = tintColors.get(prim.tintChannel());
                if (colorInt != null) {
                    r = (colorInt >> 16) & 0xFF;
                    g = (colorInt >> 8) & 0xFF;
                    b = colorInt & 0xFF;
                }
            }

            for (int idx : prim.indices()) {
                GltfSkinningEngine.skinVertex(
                    data,
                    idx,
                    jointMatrices,
                    outPos,
                    outNormal,
                    tmpPos,
                    tmpNorm
                );

                float u = texCoords[idx * 2];
                float v = texCoords[idx * 2 + 1];

                vc
                    .vertex(pose, outPos[0], outPos[1], outPos[2])
                    .color(r, g, b, 255)
                    .uv(u, 1.0f - v)
                    .overlayCoords(packedOverlay)
                    .uv2(packedLight)
                    .normal(normalMat, outNormal[0], outNormal[1], outNormal[2])
                    .endVertex();
            }
        }
    }
}